From 808de8562f39eb564da902878d650a3389b2e5ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D0=B5=D1=80=D0=B3=D0=B5=D0=B9=20=D0=9C=D0=B0=D1=80?= =?UTF-8?q?=D0=B8=D0=BD=D0=BA=D0=B5=D0=B2=D0=B8=D1=87?= Date: Fri, 5 Jun 2026 14:47:33 +0700 Subject: [PATCH] massive improve --- mattermost-dl.py | 212 ++++++++++++++++----- mm_index.py | 323 +++++++++++++++++++++++++++++++ mm_search.py | 481 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 964 insertions(+), 52 deletions(-) create mode 100644 mm_index.py create mode 100644 mm_search.py diff --git a/mattermost-dl.py b/mattermost-dl.py index 69a167a..7d15f05 100644 --- a/mattermost-dl.py +++ b/mattermost-dl.py @@ -1,13 +1,14 @@ import os import pickle import sqlite3 -from datetime import datetime, date +from datetime import datetime, date, timezone from typing import Tuple, Dict, List, Optional import getpass import requests from bs4 import BeautifulSoup from mattermostdriver import Driver +from mattermostdriver.exceptions import ResourceNotFound import pathlib import json @@ -184,6 +185,11 @@ def select_channel(d: Driver, team: str, my_user_id: str, user_id_to_name: Dict[ # The channel name consists of two user ids connected by a double underscore user_ids = channel["name"].split("__") other_user_id = user_ids[1] if user_ids[0] == my_user_id else user_ids[0] + if other_user_id not in user_id_to_name: + try: + user_id_to_name[other_user_id] = d.users.get_user(other_user_id)["username"] + except Exception: + user_id_to_name[other_user_id] = channel.get("display_name") or other_user_id channel["display_name"] = user_id_to_name[other_user_id] # Sort channels by name for easier search channels = sorted(channels, key=lambda x: x["display_name"].lower()) @@ -205,29 +211,73 @@ def select_channel(d: Driver, team: str, my_user_id: str, user_id_to_name: Dict[ return selected_channels +def _resolve_team_name(d: Driver, channel: dict, fallback_team: Optional[dict] = None) -> str: + team_id = channel.get("team_id") or "" + if team_id: + try: + return d.teams.get_team(team_id)["name"] + except ResourceNotFound: + pass + if fallback_team: + return fallback_team["name"] + return "" + + +def _parse_filter_timestamp(s: str) -> float: + for fmt in ('%Y-%m-%dT%H:%M:%SZ', '%Y-%m-%d'): + try: + dt = datetime.strptime(s, fmt) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.timestamp() + except ValueError: + continue + raise ValueError(f"Invalid date format: {s!r} (expected YYYY-MM-DD or YYYY-MM-DDTHH:MM:SSZ)") + + +def save_config(config: dict, config_filename: str = "config.json") -> None: + to_save = {k: v for k, v in config.items() if k != "password"} + with open(config_filename, "w") as f: + json.dump(to_save, f, indent=2) + + def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], output_base: str, - download_files: bool = True, before: str = None, after: str = None): + download_files: bool = True, before: str = None, after: str = None, + team: Optional[dict] = None, merge: bool = False): # Sanitize channel name channel_name = channel["display_name"].replace("\\", "").replace("/", "") print("Exporting channel", channel_name) + after_ts = _parse_filter_timestamp(after) if after else None + before_ts = _parse_filter_timestamp(before) if before else None if after: - after = datetime.strptime(after, '%Y-%m-%d').timestamp() - if before: - before = datetime.strptime(before, '%Y-%m-%d').timestamp() + print(f" fetching posts after {after}") # Get all posts for selected channel page = 0 all_posts = [] while True: print(f"Requesting channel page {page}") - posts = d.posts.get_posts_for_channel(channel["id"], params={"per_page": 200, "page": page}) + params = {"per_page": 200, "page": page} + if after_ts is not None: + # Ask server only for posts changed since last export. + # For incremental updates we typically expect few new posts, + # so a single page is enough and avoids walking the whole history. + params["since"] = int(after_ts * 1000) + posts = d.posts.get_posts_for_channel(channel["id"], params=params) if len(posts["posts"]) == 0: # If no posts are returned, we have reached the end break all_posts.extend([posts["posts"][post] for post in posts["order"]]) + + # In incremental mode (after_ts is set) we intentionally stop + # after the first non-empty page to avoid downloading the + # entire channel history again. For full dumps, keep paging. + if after_ts is not None: + break + page += 1 print(f"Found {len(all_posts)} posts") # Create output directory @@ -239,14 +289,16 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out # Filter posts by date range created = post["create_at"] / 1000 - if (before and created > before) or (after and created < after): + if (before_ts and created > before_ts) or (after_ts and created <= after_ts): continue user_id = post["user_id"] if user_id not in user_id_to_name: user_id_to_name[user_id] = d.users.get_user(user_id)["username"] username = user_id_to_name[user_id] - created = datetime.utcfromtimestamp(post["create_at"] / 1000).strftime('%Y-%m-%dT%H:%M:%SZ') + created = datetime.fromtimestamp( + post["create_at"] / 1000, timezone.utc + ).strftime('%Y-%m-%dT%H:%M:%SZ') message = post["message"] simple_post = dict(idx=i_post, id=post["id"], created=created, username=username, message=message) @@ -270,12 +322,12 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out if download_files: filename = "%03d" % i_post + "_" + file["name"] print("Downloading", file["name"]) - while True: - try: - resp = d.files.get_file(file["id"]) - break - except: - print("Downloading file failed") + try: + resp = d.files.get_file(file["id"]) + except ResourceNotFound: + print(f"File not found, skipping: {file['name']}") + filenames.append(file["name"]) + continue # Mattermost Driver unfortunately parses json files to dicts if isinstance(resp, dict): with open(output_base / filename, "w") as f: @@ -288,23 +340,42 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out simple_post["files"] = filenames simple_posts.append(simple_post) + filtered_channel_name = ''.join(filter(lambda ch: ch not in "?!/\\.;:*\"<>|", channel_name)) + output_filename = filtered_channel_name + ".json" + output_filepath = output_base / output_filename + + merged_posts = simple_posts + if merge and output_filepath.exists(): + with open(output_filepath, encoding='utf8') as f: + existing = json.load(f) + seen_ids = {post["id"] for post in existing.get("posts", [])} + merged_posts = list(existing.get("posts", [])) + new_count = 0 + for post in simple_posts: + if post["id"] not in seen_ids: + merged_posts.append(post) + seen_ids.add(post["id"]) + new_count += 1 + merged_posts.sort(key=lambda p: p["created"]) + for i, post in enumerate(merged_posts): + post["idx"] = i + print(f" merged {new_count} new posts ({len(merged_posts)} total)") + elif merge: + print(f" no existing export at '{output_filepath}', writing fresh") + output = { "channel": { "name": channel["name"], "display_name": channel["display_name"], "header": channel["header"], "id": channel["id"], - "team": d.teams.get_team(channel["team_id"])["name"], - "team_id": channel["team_id"], - "exported_at": datetime.now().strftime('%Y-%m-%dT%H:%M:%SZ') + "team": _resolve_team_name(d, channel, team), + "team_id": channel.get("team_id") or "", + "exported_at": datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ') }, - "posts": simple_posts + "posts": merged_posts } - # Export posts to json file - filtered_channel_name = ''.join(filter(lambda ch: ch not in "?!/\\.;:*\"<>|", channel_name)) - output_filename = filtered_channel_name + ".json" - output_filepath = output_base / output_filename with open(output_filepath, "w", encoding='utf8') as f: json.dump(output, f, indent=2, ensure_ascii=False) print(f"Exported channel data to '{output_filepath}'") @@ -350,29 +421,34 @@ def complete_config(config: dict, config_filename: str = "config.json") -> dict: password = getpass.getpass("Enter your password (hidden): ") elif config["login_mode"] == "gitlab": - if config.get("gitlab_url", False): - print(f"Using GitLab URL '{config['gitlab_url']}' from config") + # If we already have a token in config, reuse it and do not ask + # for GitLab credentials again. + if config.get("token"): + print("Using existing token from config for GitLab login_mode.") else: - config["gitlab_url"] = input( - f"Please input GitLab URL [{GITLAB_URL_DEFAULT}]: " - ) or GITLAB_URL_DEFAULT - config_changed = True + if config.get("gitlab_url", False): + print(f"Using GitLab URL '{config['gitlab_url']}' from config") + else: + config["gitlab_url"] = input( + f"Please input GitLab URL [{GITLAB_URL_DEFAULT}]: " + ) or GITLAB_URL_DEFAULT + config_changed = True - if config.get("username", False): - print(f"Using username '{config['username']}' from config") - else: - config["username"] = input("Please input your GitLab username: ") - config_changed = True + if config.get("username", False): + print(f"Using username '{config['username']}' from config") + else: + config["username"] = input("Please input your GitLab username: ") + config_changed = True - password = getpass.getpass("Enter your GitLab password (hidden): ") - token = get_token_via_gitlab( - config["host"], config["gitlab_url"], config["username"], password, - config.get("cookie_file", COOKIE_FILE_DEFAULT) - ) - if not token: - raise SystemExit("GitLab authentication failed.") - config["token"] = token - config_changed = True + password = getpass.getpass("Enter your GitLab password (hidden): ") + token = get_token_via_gitlab( + config["host"], config["gitlab_url"], config["username"], password, + config.get("cookie_file", COOKIE_FILE_DEFAULT) + ) + if not token: + raise SystemExit("GitLab authentication failed.") + config["token"] = token + config_changed = True else: if config.get("token", False): print(f"Using token '{config['token']}' from config") @@ -465,23 +541,55 @@ def find_mmauthtoken_firefox(host): if __name__ == '__main__': - config = get_config_from_json() - config = complete_config(config) + config_filename = "config.json" + config = get_config_from_json(config_filename) + config = complete_config(config, config_filename) + + incremental = config.get("incremental", bool(config.get("last_export_at"))) + if incremental and config.get("last_export_at"): + print(f"Incremental update from {config['last_export_at']}") + output_base = config.get("output_base") or ("results/" + date.today().strftime("%Y%m%d")) + after = config["last_export_at"] + before = config.get("before") + merge = True + else: + output_base = "results/" + date.today().strftime("%Y%m%d") + after = config.get("after") + before = config.get("before") + merge = False - output_base = "results/" + date.today().strftime("%Y%m%d") print(f"Storing downloaded data in {output_base}") + export_started_at = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ') - # Range of posts to be exported as string in format "YYYY-MM-DD". Use None if no filter should be applied - after = config.get("after", None) - before = config.get("before", None) - - d = connect(config["host"], config.get("token", None), - config.get("username", None), config.get("password", None)) + try: + d = connect( + config["host"], + config.get("token", None), + config.get("username", None), + config.get("password", None), + ) + except Exception as e: + # If token-based authentication failed, drop the token and ask + # the user to rerun the script so that fresh credentials can be + # obtained or entered. + if config.get("token"): + print("Token-based authentication failed, removing stored token.") + config.pop("token", None) + save_config(config, config_filename) + print("Please restart the script to reconfigure authentication.") + raise SystemExit(1) from e + raise user_id_to_name, my_user_id = get_users(d) team = select_team(d, my_user_id) channels = select_channel(d, team, my_user_id, user_id_to_name) for i_channel, channel in enumerate(channels): print(f"Start export of channel {i_channel + 1}/{len(channels)}") export_channel(d, channel, user_id_to_name, output_base, config["download_files"], - before, after) + before, after, team=team, merge=merge) + + config["last_export_at"] = export_started_at + config["output_base"] = output_base + config["incremental"] = True + save_config(config, config_filename) + print(f"Updated last_export_at to {export_started_at}") print("Finished export") diff --git a/mm_index.py b/mm_index.py new file mode 100644 index 0000000..fdd6f34 --- /dev/null +++ b/mm_index.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +import argparse +import datetime as dt +import json +import os +import sqlite3 +from typing import Any, Dict, Iterable, List, Optional, Tuple + + +def iter_result_files(root: str) -> Iterable[str]: + for dirpath, _dirnames, filenames in os.walk(root): + for name in filenames: + if name.endswith(".json"): + yield os.path.join(dirpath, name) + + +def parse_post_created(s: str) -> Optional[str]: + try: + created = dt.datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc) + except ValueError: + return None + return created.isoformat() + + +def parse_exported_at(channel_obj: Dict[str, Any]) -> Optional[str]: + raw = channel_obj.get("exported_at") + if not isinstance(raw, str): + return None + try: + exported = dt.datetime.strptime(raw, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc) + except ValueError: + return None + return exported.isoformat() + + +def connect_db(db_path: str) -> sqlite3.Connection: + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS sources ( + path TEXT PRIMARY KEY, + mtime_ns INTEGER NOT NULL, + size INTEGER NOT NULL, + indexed_at TEXT NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS posts ( + post_id TEXT PRIMARY KEY, + idx INTEGER, + created TEXT NOT NULL, + username TEXT NOT NULL, + message TEXT NOT NULL, + files_json TEXT, + channel_id TEXT, + channel_name TEXT, + channel_display_name TEXT, + team TEXT, + source_path TEXT NOT NULL, + exported_at TEXT + ) + """ + ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_created ON posts(created)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_username ON posts(username)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_channel_display ON posts(channel_display_name)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_team ON posts(team)") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS attachments ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + post_id TEXT NOT NULL, + filename TEXT NOT NULL, + filepath TEXT NOT NULL, + is_text INTEGER NOT NULL, + FOREIGN KEY(post_id) REFERENCES posts(post_id) ON DELETE CASCADE + ) + """ + ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_attachments_post_id ON attachments(post_id)") + conn.execute( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS posts_fts + USING fts5( + post_id UNINDEXED, + message, + tokenize='unicode61' + ) + """ + ) + conn.execute( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS attachments_fts + USING fts5( + attachment_id UNINDEXED, + post_id UNINDEXED, + filename, + content, + tokenize='unicode61' + ) + """ + ) + return conn + + +def get_source_state(conn: sqlite3.Connection, path: str) -> Optional[Tuple[int, int]]: + row = conn.execute("SELECT mtime_ns, size FROM sources WHERE path = ?", (path,)).fetchone() + if not row: + return None + return int(row["mtime_ns"]), int(row["size"]) + + +def source_changed(path: str, known_state: Optional[Tuple[int, int]]) -> bool: + st = os.stat(path) + if known_state is None: + return True + return known_state != (st.st_mtime_ns, st.st_size) + + +def load_posts(path: str) -> Optional[Tuple[Dict[str, Any], List[Dict[str, Any]]]]: + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + return None + + channel_obj = data.get("channel") + posts = data.get("posts") + if not isinstance(channel_obj, dict) or not isinstance(posts, list): + return None + return channel_obj, posts + + +def upsert_file(conn: sqlite3.Connection, path: str) -> Tuple[int, int]: + loaded = load_posts(path) + if loaded is None: + return 0, 0 + + channel_obj, posts = loaded + channel_id = str(channel_obj.get("id") or "") + channel_name = str(channel_obj.get("name") or "") + channel_display_name = str(channel_obj.get("display_name") or "") + team = str(channel_obj.get("team") or "") + exported_at = parse_exported_at(channel_obj) + + # Remove any existing attachments for posts coming from this source file + conn.execute( + """ + DELETE FROM attachments + WHERE post_id IN (SELECT post_id FROM posts WHERE source_path = ?) + """, + (path,), + ) + conn.execute( + """ + DELETE FROM attachments_fts + WHERE post_id IN (SELECT post_id FROM posts WHERE source_path = ?) + """, + (path,), + ) + + base_dir = os.path.dirname(path) + max_text_size = 2 * 1024 * 1024 # 2 MiB + + inserted_or_updated = 0 + skipped = 0 + for p in posts: + post_id = p.get("id") + created_raw = p.get("created") + username = p.get("username") + message = p.get("message") + if not isinstance(post_id, str) or not post_id: + skipped += 1 + continue + if not isinstance(created_raw, str) or not isinstance(username, str): + skipped += 1 + continue + if not isinstance(message, str): + message = "" + + created = parse_post_created(created_raw) + if created is None: + skipped += 1 + continue + + files = p.get("files") + files_json = json.dumps(files, ensure_ascii=False) if isinstance(files, list) else None + idx = p.get("idx") if isinstance(p.get("idx"), int) else None + + conn.execute( + """ + INSERT INTO posts ( + post_id, idx, created, username, message, files_json, + channel_id, channel_name, channel_display_name, team, + source_path, exported_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(post_id) DO UPDATE SET + idx = excluded.idx, + created = excluded.created, + username = excluded.username, + message = excluded.message, + files_json = excluded.files_json, + channel_id = excluded.channel_id, + channel_name = excluded.channel_name, + channel_display_name = excluded.channel_display_name, + team = excluded.team, + source_path = excluded.source_path, + exported_at = excluded.exported_at + """, + ( + post_id, + idx, + created, + username, + message, + files_json, + channel_id, + channel_name, + channel_display_name, + team, + path, + exported_at, + ), + ) + conn.execute("DELETE FROM posts_fts WHERE post_id = ?", (post_id,)) + conn.execute("INSERT INTO posts_fts(post_id, message) VALUES (?, ?)", (post_id, message)) + + # Index attachments (file names and, for text files, contents) + if isinstance(files, list) and files: + for filename in files: + if not isinstance(filename, str) or not filename: + continue + disk_name = f"{idx:03d}_{filename}" if isinstance(idx, int) else filename + disk_path = os.path.join(base_dir, disk_name) + is_text = 0 + content: Optional[str] = None + try: + st = os.stat(disk_path) + if st.st_size <= max_text_size: + with open(disk_path, "r", encoding="utf-8") as f: + content = f.read() + is_text = 1 + except (OSError, UnicodeDecodeError): + is_text = 0 + content = None + + cur = conn.execute( + """ + INSERT INTO attachments (post_id, filename, filepath, is_text) + VALUES (?, ?, ?, ?) + """, + (post_id, filename, disk_path, is_text), + ) + attachment_id = cur.lastrowid + conn.execute( + """ + INSERT INTO attachments_fts (attachment_id, post_id, filename, content) + VALUES (?, ?, ?, ?) + """, + (attachment_id, post_id, filename, content if is_text else None), + ) + inserted_or_updated += 1 + + st = os.stat(path) + conn.execute( + """ + INSERT INTO sources(path, mtime_ns, size, indexed_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(path) DO UPDATE SET + mtime_ns = excluded.mtime_ns, + size = excluded.size, + indexed_at = excluded.indexed_at + """, + (path, st.st_mtime_ns, st.st_size, dt.datetime.now(dt.timezone.utc).isoformat()), + ) + return inserted_or_updated, skipped + + +def main() -> int: + ap = argparse.ArgumentParser(description="Build/update sqlite index for Mattermost export JSON dumps.") + ap.add_argument("-r", "--root", default="results", help="Root folder containing dump directories.") + ap.add_argument("--db", default="mattermost_index.sqlite", help="SQLite index database file.") + ap.add_argument("--reindex-all", action="store_true", help="Re-index all files, ignoring sources cache.") + args = ap.parse_args() + + conn = connect_db(args.db) + files = list(iter_result_files(args.root)) + + processed_files = 0 + skipped_files = 0 + rows_upserted = 0 + rows_skipped = 0 + + with conn: + for path in files: + known_state = None if args.reindex_all else get_source_state(conn, path) + if not args.reindex_all and not source_changed(path, known_state): + skipped_files += 1 + continue + processed_files += 1 + upserted, skipped = upsert_file(conn, path) + rows_upserted += upserted + rows_skipped += skipped + + total_posts = conn.execute("SELECT COUNT(*) AS c FROM posts").fetchone()["c"] + total_sources = conn.execute("SELECT COUNT(*) AS c FROM sources").fetchone()["c"] + conn.close() + + print(f"DB: {args.db}") + print(f"Files scanned: {len(files)} (processed={processed_files}, unchanged_skipped={skipped_files})") + print(f"Rows upserted: {rows_upserted}, invalid_posts_skipped: {rows_skipped}") + print(f"Totals: posts={total_posts}, sources={total_sources}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/mm_search.py b/mm_search.py new file mode 100644 index 0000000..16207b0 --- /dev/null +++ b/mm_search.py @@ -0,0 +1,481 @@ +#!/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()) +