massive improve
This commit is contained in:
+160
-52
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user