Сергей Маринкевич
482 lines
15 KiB
Python
482 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import datetime as dt
|
|
import json
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
from typing import Any, Dict, Iterable, List, Optional
|
|
|
|
|
|
def iter_result_files(root: str) -> Iterable[str]:
|
|
for dirpath, _dirnames, filenames in os.walk(root):
|
|
for name in filenames:
|
|
if not name.endswith(".json"):
|
|
continue
|
|
yield os.path.join(dirpath, name)
|
|
|
|
|
|
def parse_iso_date(s: str) -> dt.datetime:
|
|
# Accept YYYY-MM-DD or full ISO with time; normalize to UTC
|
|
try:
|
|
if "T" in s:
|
|
return dt.datetime.fromisoformat(s.replace("Z", "+00:00")).astimezone(dt.timezone.utc)
|
|
return dt.datetime.strptime(s, "%Y-%m-%d").replace(tzinfo=dt.timezone.utc)
|
|
except ValueError as e:
|
|
raise argparse.ArgumentTypeError(f"invalid date '{s}': {e}") from e
|
|
|
|
|
|
def parse_post_created(s: str) -> dt.datetime:
|
|
# Exporter uses '%Y-%m-%dT%H:%M:%SZ'
|
|
return dt.datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc)
|
|
|
|
|
|
def make_matcher(pattern: str, use_regex: bool, ignore_case: bool):
|
|
if use_regex:
|
|
flags = re.IGNORECASE if ignore_case else 0
|
|
rx = re.compile(pattern, flags)
|
|
|
|
def match(text: str) -> bool:
|
|
return bool(rx.search(text))
|
|
|
|
return match
|
|
|
|
needle = pattern.lower() if ignore_case else pattern
|
|
|
|
def match(text: str) -> bool:
|
|
hay = text.lower() if ignore_case else text
|
|
return needle in hay
|
|
|
|
return match
|
|
|
|
|
|
def build_fts_query(query: str) -> Optional[str]:
|
|
# Tokenize user query and run prefix search by tokens.
|
|
tokens = re.findall(r"\w+", query, flags=re.UNICODE)
|
|
if not tokens:
|
|
return None
|
|
return " AND ".join(f"{tok}*" for tok in tokens)
|
|
|
|
|
|
def search_index(
|
|
db_path: str,
|
|
query: str,
|
|
user: Optional[str],
|
|
channel: Optional[str],
|
|
team: Optional[str],
|
|
after: Optional[dt.datetime],
|
|
before: Optional[dt.datetime],
|
|
limit: Optional[int],
|
|
ignore_case: bool,
|
|
) -> Iterable[Dict[str, Any]]:
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
filters = []
|
|
params: List[Any] = []
|
|
|
|
if user:
|
|
filters.append("p.username = ?")
|
|
params.append(user)
|
|
if channel:
|
|
filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
|
|
params.extend([channel, channel])
|
|
if team:
|
|
filters.append("p.team = ?")
|
|
params.append(team)
|
|
if after:
|
|
filters.append("p.created >= ?")
|
|
params.append(after.isoformat())
|
|
if before:
|
|
filters.append("p.created <= ?")
|
|
params.append(before.isoformat())
|
|
|
|
where_filters = f"AND {' AND '.join(filters)}" if filters else ""
|
|
fts_q = build_fts_query(query)
|
|
if fts_q:
|
|
sql = f"""
|
|
SELECT
|
|
p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
|
|
p.channel_display_name, p.channel_name, p.team, p.source_path
|
|
FROM posts p
|
|
JOIN posts_fts fts ON fts.post_id = p.post_id
|
|
WHERE fts.message MATCH ?
|
|
{where_filters}
|
|
ORDER BY p.created ASC
|
|
"""
|
|
params = [fts_q] + params
|
|
else:
|
|
cmp_col = "LOWER(p.message)" if ignore_case else "p.message"
|
|
needle = query.lower() if ignore_case else query
|
|
sql = f"""
|
|
SELECT
|
|
p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
|
|
p.channel_display_name, p.channel_name, p.team, p.source_path
|
|
FROM posts p
|
|
WHERE {cmp_col} LIKE ?
|
|
{where_filters}
|
|
ORDER BY p.created ASC
|
|
"""
|
|
params = [f"%{needle}%"] + params
|
|
|
|
if limit is not None:
|
|
sql += " LIMIT ?"
|
|
params.append(limit)
|
|
|
|
rows = conn.execute(sql, params).fetchall()
|
|
for row in rows:
|
|
files: Optional[List[str]] = None
|
|
if row["files_json"]:
|
|
try:
|
|
maybe_files = json.loads(row["files_json"])
|
|
if isinstance(maybe_files, list):
|
|
files = maybe_files
|
|
except json.JSONDecodeError:
|
|
files = None
|
|
post: Dict[str, Any] = {
|
|
"id": row["post_id"],
|
|
"idx": row["idx"],
|
|
"created": row["created"].replace("+00:00", "Z"),
|
|
"username": row["username"],
|
|
"message": row["message"],
|
|
}
|
|
if files:
|
|
post["files"] = files
|
|
|
|
yield {
|
|
"file": row["source_path"],
|
|
"channel": row["channel_display_name"] or row["channel_name"],
|
|
"team": row["team"],
|
|
"post": post,
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def search_attachments_index(
|
|
db_path: str,
|
|
query: str,
|
|
user: Optional[str],
|
|
channel: Optional[str],
|
|
team: Optional[str],
|
|
after: Optional[dt.datetime],
|
|
before: Optional[dt.datetime],
|
|
limit: Optional[int],
|
|
mode: str,
|
|
) -> Iterable[Dict[str, Any]]:
|
|
"""
|
|
Search in attachments index.
|
|
|
|
mode: "name" to search by file name, "content" to search by text content.
|
|
"""
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
filters = []
|
|
params: List[Any] = []
|
|
|
|
if user:
|
|
filters.append("p.username = ?")
|
|
params.append(user)
|
|
if channel:
|
|
filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
|
|
params.extend([channel, channel])
|
|
if team:
|
|
filters.append("p.team = ?")
|
|
params.append(team)
|
|
if after:
|
|
filters.append("p.created >= ?")
|
|
params.append(after.isoformat())
|
|
if before:
|
|
filters.append("p.created <= ?")
|
|
params.append(before.isoformat())
|
|
|
|
where_filters = f"AND {' AND '.join(filters)}" if filters else ""
|
|
|
|
fts_column = "filename" if mode == "name" else "content"
|
|
fts_q = build_fts_query(query)
|
|
if not fts_q:
|
|
return []
|
|
|
|
sql = f"""
|
|
SELECT
|
|
p.post_id,
|
|
p.idx,
|
|
p.created,
|
|
p.username,
|
|
p.message,
|
|
p.files_json,
|
|
p.channel_display_name,
|
|
p.channel_name,
|
|
p.team,
|
|
p.source_path,
|
|
a.filename AS attachment_filename,
|
|
a.filepath AS attachment_path,
|
|
a.is_text AS attachment_is_text
|
|
FROM attachments a
|
|
JOIN attachments_fts aft ON aft.attachment_id = a.id
|
|
JOIN posts p ON p.post_id = a.post_id
|
|
WHERE aft.{fts_column} MATCH ?
|
|
{where_filters}
|
|
ORDER BY p.created ASC
|
|
"""
|
|
|
|
params = [fts_q] + params
|
|
if limit is not None:
|
|
sql += " LIMIT ?"
|
|
params.append(limit)
|
|
|
|
rows = conn.execute(sql, params).fetchall()
|
|
for row in rows:
|
|
files: Optional[List[str]] = None
|
|
if row["files_json"]:
|
|
try:
|
|
maybe_files = json.loads(row["files_json"])
|
|
if isinstance(maybe_files, list):
|
|
files = maybe_files
|
|
except json.JSONDecodeError:
|
|
files = None
|
|
post: Dict[str, Any] = {
|
|
"id": row["post_id"],
|
|
"idx": row["idx"],
|
|
"created": row["created"].replace("+00:00", "Z"),
|
|
"username": row["username"],
|
|
"message": row["message"],
|
|
}
|
|
if files:
|
|
post["files"] = files
|
|
|
|
attachment = {
|
|
"filename": row["attachment_filename"],
|
|
"path": row["attachment_path"],
|
|
"is_text": bool(row["attachment_is_text"]),
|
|
}
|
|
|
|
yield {
|
|
"file": row["source_path"],
|
|
"channel": row["channel_display_name"] or row["channel_name"],
|
|
"team": row["team"],
|
|
"post": post,
|
|
"attachment": attachment,
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def search_file(
|
|
path: str,
|
|
matcher,
|
|
user: Optional[str],
|
|
channel: Optional[str],
|
|
team: Optional[str],
|
|
after: Optional[dt.datetime],
|
|
before: Optional[dt.datetime],
|
|
) -> Iterable[Dict[str, Any]]:
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except (OSError, json.JSONDecodeError):
|
|
return []
|
|
|
|
ch = data.get("channel", {}) or {}
|
|
posts: List[Dict[str, Any]] = data.get("posts") or []
|
|
|
|
channel_name = ch.get("display_name") or ch.get("name") or ""
|
|
team_name = ch.get("team") or ""
|
|
|
|
if channel and channel not in (channel_name, ch.get("name", "")):
|
|
return []
|
|
if team and team != team_name:
|
|
return []
|
|
|
|
for p in posts:
|
|
msg = p.get("message") or ""
|
|
username = p.get("username") or ""
|
|
created_raw = p.get("created")
|
|
if not isinstance(created_raw, str):
|
|
continue
|
|
try:
|
|
created = parse_post_created(created_raw)
|
|
except ValueError:
|
|
continue
|
|
|
|
if user and username != user:
|
|
continue
|
|
if after and created < after:
|
|
continue
|
|
if before and created > before:
|
|
continue
|
|
if not matcher(msg):
|
|
continue
|
|
|
|
yield {
|
|
"file": path,
|
|
"channel": channel_name,
|
|
"team": team_name,
|
|
"post": p,
|
|
}
|
|
|
|
|
|
def print_human(matches: Iterable[Dict[str, Any]], limit: Optional[int], show_file: bool) -> None:
|
|
count = 0
|
|
for m in matches:
|
|
if limit is not None and count >= limit:
|
|
break
|
|
p = m["post"]
|
|
created = p.get("created", "")
|
|
username = p.get("username", "")
|
|
msg = (p.get("message") or "").replace("\n", "\\n")
|
|
attachment = m.get("attachment")
|
|
attachment_suffix = ""
|
|
if isinstance(attachment, dict) and attachment.get("filename"):
|
|
attachment_suffix = f" [file: {attachment['filename']}]"
|
|
if show_file:
|
|
file_rel = os.path.relpath(m["file"])
|
|
idx = p.get("idx", "?")
|
|
prefix = f"{file_rel}:{idx}: "
|
|
else:
|
|
prefix = ""
|
|
channel = m.get("channel") or "?"
|
|
team = m.get("team") or ""
|
|
team_prefix = f"[{team}]" if team else ""
|
|
print(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
|
|
count += 1
|
|
|
|
|
|
def print_json(matches: Iterable[Dict[str, Any]], limit: Optional[int]) -> None:
|
|
out: List[Dict[str, Any]] = []
|
|
count = 0
|
|
for m in matches:
|
|
if limit is not None and count >= limit:
|
|
break
|
|
out.append(m)
|
|
count += 1
|
|
json.dump(out, fp=os.sys.stdout, ensure_ascii=False, indent=2)
|
|
print()
|
|
|
|
|
|
def main(argv: Optional[List[str]] = None) -> int:
|
|
ap = argparse.ArgumentParser(
|
|
description="Search through Mattermost export JSON files produced by mattermost-dl.py.",
|
|
)
|
|
ap.add_argument("query", help="Search string or regex (see --regex).")
|
|
ap.add_argument(
|
|
"-r",
|
|
"--root",
|
|
default="results",
|
|
help="Root directory with export results (default: results).",
|
|
)
|
|
ap.add_argument("--user", help="Filter by username (exact match).")
|
|
ap.add_argument("--channel", help="Filter by channel display_name or name (exact match).")
|
|
ap.add_argument("--team", help="Filter by team name (exact match).")
|
|
ap.add_argument("--after", type=parse_iso_date, help="Only messages on/after this date (YYYY-MM-DD or ISO).")
|
|
ap.add_argument("--before", type=parse_iso_date, help="Only messages on/before this date (YYYY-MM-DD or ISO).")
|
|
ap.add_argument("--regex", action="store_true", help="Treat query as regular expression.")
|
|
group = ap.add_mutually_exclusive_group()
|
|
group.add_argument(
|
|
"--file-name",
|
|
action="store_true",
|
|
help="Search by attachment file name instead of message text (requires sqlite index).",
|
|
)
|
|
group.add_argument(
|
|
"--file-content",
|
|
action="store_true",
|
|
help="Search by attachment text content instead of message text (requires sqlite index; only UTF-8 text files are indexed).",
|
|
)
|
|
ap.add_argument(
|
|
"-i",
|
|
"--ignore-case",
|
|
action="store_true",
|
|
help="Case-insensitive search (default for substring search).",
|
|
)
|
|
ap.add_argument("--json", action="store_true", help="Output matches as JSON array (for jq, etc.).")
|
|
ap.add_argument(
|
|
"-n",
|
|
"--limit",
|
|
type=int,
|
|
help="Maximum number of matches to output.",
|
|
)
|
|
ap.add_argument("--db", default="mattermost_index.sqlite", help="SQLite index database path.")
|
|
ap.add_argument(
|
|
"--no-index",
|
|
action="store_true",
|
|
help="Do not use sqlite index even if DB exists.",
|
|
)
|
|
ap.add_argument(
|
|
"--show-file",
|
|
action="store_true",
|
|
help="Include source filename and per-file index in human output.",
|
|
)
|
|
|
|
args = ap.parse_args(argv)
|
|
|
|
# Default ignore-case for non-regex search unless explicitly disabled
|
|
ignore_case = args.ignore_case or not args.regex
|
|
matcher = make_matcher(args.query, args.regex, ignore_case)
|
|
|
|
if (args.file_name or args.file_content) and args.regex:
|
|
raise SystemExit("--file-name/--file-content do not support --regex yet.")
|
|
|
|
use_index = (not args.no_index) and os.path.exists(args.db)
|
|
if args.file_name or args.file_content:
|
|
if not use_index:
|
|
raise SystemExit(
|
|
"Attachment search (--file-name/--file-content) requires sqlite index. "
|
|
"Create it with mm_index.py and do not use --no-index."
|
|
)
|
|
mode = "name" if args.file_name else "content"
|
|
matches_iter = search_attachments_index(
|
|
db_path=args.db,
|
|
query=args.query,
|
|
user=args.user,
|
|
channel=args.channel,
|
|
team=args.team,
|
|
after=args.after,
|
|
before=args.before,
|
|
limit=args.limit,
|
|
mode=mode,
|
|
)
|
|
output_limit = None
|
|
elif use_index and not args.regex:
|
|
matches_iter = search_index(
|
|
db_path=args.db,
|
|
query=args.query,
|
|
user=args.user,
|
|
channel=args.channel,
|
|
team=args.team,
|
|
after=args.after,
|
|
before=args.before,
|
|
limit=args.limit,
|
|
ignore_case=ignore_case,
|
|
)
|
|
# Limit already applied in SQL.
|
|
output_limit = None
|
|
else:
|
|
files = list(iter_result_files(args.root))
|
|
matches_iter = (
|
|
m
|
|
for path in files
|
|
for m in search_file(
|
|
path=path,
|
|
matcher=matcher,
|
|
user=args.user,
|
|
channel=args.channel,
|
|
team=args.team,
|
|
after=args.after,
|
|
before=args.before,
|
|
)
|
|
)
|
|
output_limit = args.limit
|
|
|
|
if args.json:
|
|
print_json(matches_iter, output_limit)
|
|
else:
|
|
print_human(matches_iter, output_limit, show_file=args.show_file)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|
|
|