massive improve

This commit is contained in:
Сергей Маринкевич
2026-06-05 14:47:33 +07:00
parent 3b332d2e0c
commit 808de8562f
3 changed files with 964 additions and 52 deletions
+160 -52
View File
@@ -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")