346 lines
16 KiB
Python
346 lines
16 KiB
Python
import sys
|
|
import asyncio
|
|
import discord
|
|
import logging
|
|
import time
|
|
import re
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from rich.console import Console
|
|
from rich.prompt import Prompt, Confirm
|
|
from rich.table import Table
|
|
from rich.panel import Panel
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
|
|
|
|
from src.core.configuration import load_config, save_config
|
|
from src.core.base import MigrationContext
|
|
from src.disco_reaper.exporter import DiscordExporter
|
|
|
|
console = Console()
|
|
|
|
class DiscoReaperCLI:
|
|
"""CLI app to manage Discord server data export."""
|
|
|
|
def __init__(self, config_path="config.yaml"):
|
|
self.config_path = config_path
|
|
try:
|
|
self.config = load_config(self.config_path)
|
|
# We only need the Discord side of MigrationContext
|
|
self.engine = MigrationContext(self.config, target_platform="fluxer")
|
|
self.exporter = DiscordExporter(self.engine.discord_reader)
|
|
except Exception as e:
|
|
console.print(f"[bold red]Failed to load config: {e}[/bold red]")
|
|
sys.exit(1)
|
|
|
|
self.tokens_valid = False
|
|
|
|
async def validate_config(self):
|
|
self.validation_results = {
|
|
"discord_token": False, "discord_bot_name": None,
|
|
"discord_server": False, "discord_server_name": None,
|
|
"discord_timeout": False
|
|
}
|
|
|
|
d_token = self.config.discord_bot_token
|
|
fillers = ["DISCORD_BOT_TOKEN", "000000000000000000", "DISCORD_SERVER_ID", "", None]
|
|
discord_dummy = d_token in fillers or self.config.discord_server_id in fillers
|
|
|
|
if discord_dummy:
|
|
console.print("[bold yellow]Discord setup incomplete in config.yaml[/bold yellow]")
|
|
return False
|
|
|
|
try:
|
|
with console.status("[yellow]Validating Discord token...[/yellow]"):
|
|
res = await self.engine.discord_reader.validate()
|
|
self.validation_results["discord_token"] = res.get("token", False)
|
|
self.validation_results["discord_bot_name"] = res.get("bot_name")
|
|
self.validation_results["discord_server"] = res.get("server", False)
|
|
self.validation_results["discord_server_name"] = res.get("server_name")
|
|
|
|
if not res.get("token"):
|
|
console.print("[bold red]Discord Token validation failed.[/bold red]")
|
|
elif not res.get("server"):
|
|
console.print("[bold red]Discord Server ID validation failed.[/bold red]")
|
|
else:
|
|
self.tokens_valid = True
|
|
return True
|
|
except Exception as e:
|
|
console.print(f"[bold red]Discord validation failed: {e}[/bold red]")
|
|
|
|
return False
|
|
|
|
def get_backup_info(self):
|
|
"""Checks for existing backup and returns formatted timestamp."""
|
|
d_name = self.validation_results.get("discord_server_name")
|
|
d_id = self.config.discord_server_id
|
|
if not d_name or not d_id or d_id == "DISCORD_SERVER_ID":
|
|
return None
|
|
|
|
safe_name = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', d_name)
|
|
export_path = Path(".") / f"EXPORT-{safe_name}-{d_id}"
|
|
profile_file = export_path / "server_profile.json"
|
|
|
|
if profile_file.exists():
|
|
try:
|
|
with open(profile_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
ts_str = data.get("last_backup")
|
|
if ts_str:
|
|
dt = datetime.fromisoformat(ts_str)
|
|
return dt.strftime("%d-%b-%Y %H:%M")
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
async def run(self):
|
|
await self.validate_config()
|
|
|
|
while True:
|
|
console.print("")
|
|
console.print(Panel.fit("Disco Reaper - Server Backup Tool", style="bold green"))
|
|
|
|
d_name = self.validation_results.get("discord_server_name")
|
|
d_display = f"[bold green]\"{d_name}\"[/bold green]" if d_name else "[bold red]NOT CONNECTED[/bold red]"
|
|
console.print(f"[bold cyan]Source Server:[/bold cyan] {d_display}")
|
|
|
|
b_name = self.validation_results.get("discord_bot_name")
|
|
b_display = f"[bold green]\"{b_name}\"[/bold green]" if b_name else "[bold red]UNKNOWN[/bold red]"
|
|
console.print(f"[bold cyan]Bot name:[/bold cyan] {b_display}")
|
|
|
|
backup_ts = self.get_backup_info()
|
|
if backup_ts:
|
|
console.print(f"[bold cyan]Backup Found:[/bold cyan] [bold yellow]{backup_ts}[/bold yellow]")
|
|
|
|
console.print("\n[bold]Main Menu[/bold]")
|
|
console.print("(1) Backup Server Profile")
|
|
console.print("(2) Backup Channel Messages")
|
|
console.print("(3) Update & Sync Backup")
|
|
console.print("(4) Configuration")
|
|
console.print("(Q) Exit")
|
|
|
|
choice = Prompt.ask("\nSelect an option", choices=["1", "2", "3", "4", "Q"], default="Q", show_choices=False).upper()
|
|
|
|
if choice == "1":
|
|
await self.backup_server_profile()
|
|
elif choice == "2":
|
|
await self.backup_messages()
|
|
elif choice == "3":
|
|
await self.sync_backup()
|
|
elif choice == "4":
|
|
await self.edit_configuration()
|
|
elif choice == "Q":
|
|
await self.engine.close_connections()
|
|
break
|
|
|
|
async def backup_server_profile(self):
|
|
"""Option 1: Backup name, banner, logo, roles, structure, and custom assets."""
|
|
if not self.tokens_valid: return
|
|
try:
|
|
await self.engine.discord_reader.start()
|
|
meta = await self.exporter.setup()
|
|
|
|
with console.status("[yellow]Backing up server profile & skeleton...[/yellow]"):
|
|
# 1. Profile & Assets
|
|
await self.exporter.export_metadata()
|
|
await self.exporter.download_server_assets()
|
|
|
|
# 2. Roles
|
|
try:
|
|
role_data = await self.exporter.export_roles()
|
|
r_count = len(role_data)
|
|
except discord.Forbidden:
|
|
r_count = 0
|
|
console.print("[yellow]Warning: Could not backup roles (Missing Access).[/yellow]")
|
|
|
|
# 3. Structure
|
|
_, cat_count, chan_count = await self.exporter.export_channels_structure()
|
|
|
|
# 4. Emojis & Stickers
|
|
e_count, s_count = await self.exporter.export_assets()
|
|
|
|
console.print(f"[bold green]Server Profile backed up to: {self.exporter.export_path}[/bold green]")
|
|
console.print("\n[bold]Profile Backup Stats:[/bold]")
|
|
console.print("- name, logo & banner")
|
|
console.print(f"- {cat_count} categories & {chan_count} channels")
|
|
console.print(f"- {r_count} roles")
|
|
console.print(f"- {e_count} emojis, {s_count} stickers.")
|
|
except discord.Forbidden as e:
|
|
console.print(f"[bold red]Backup failed: {e}[/bold red]")
|
|
console.print("[yellow]Hint: Missing permissions. Check bot 'Guilds' and 'Channel' access.[/yellow]")
|
|
except Exception as e:
|
|
console.print(f"[bold red]Backup failed: {e}[/bold red]")
|
|
finally:
|
|
await self.engine.close_connections()
|
|
|
|
async def backup_messages(self, force_overwrite=None, skip_found_prompt=False, auto_sync_existing=False):
|
|
"""Option 2: Backup Channels structure and all message history."""
|
|
if not self.tokens_valid: return
|
|
try:
|
|
await self.engine.discord_reader.start()
|
|
await self.exporter.setup()
|
|
|
|
# 1. Export structure first
|
|
with console.status("[yellow]Exporting channel structure...[/yellow]"):
|
|
await self.exporter.export_channels_structure()
|
|
|
|
# 2. Select Channels
|
|
with console.status("[yellow]Fetching channels & categories...[/yellow]"):
|
|
all_channels = await self.engine.discord_reader.get_channels()
|
|
all_categories = await self.engine.discord_reader.get_categories()
|
|
cat_map = {c.id: c.name for c in all_categories}
|
|
|
|
# Filter for exportable channels
|
|
eligible_channels = [
|
|
c for c in all_channels
|
|
if c.type in [discord.ChannelType.text, discord.ChannelType.news]
|
|
]
|
|
|
|
if not eligible_channels:
|
|
console.print("[yellow]No text/news channels found to backup.[/yellow]")
|
|
return
|
|
|
|
selected_channels = []
|
|
|
|
if auto_sync_existing:
|
|
selected_channels = [
|
|
c for c in eligible_channels
|
|
if (self.exporter.export_path / "message_backup" / f"{c.id}.json").exists()
|
|
]
|
|
if not selected_channels:
|
|
console.print("[yellow]No existing backups found to sync.[/yellow]")
|
|
return
|
|
# In auto-sync mode, we always use sync (force_overwrite=False) and skip prompts
|
|
force_overwrite = False
|
|
skip_found_prompt = True
|
|
else:
|
|
console.print(f"\n[bold]Select Channels to Backup ({len(eligible_channels)} total):[/bold]")
|
|
|
|
# Identify duplicate names to show categories for them
|
|
name_counts = {}
|
|
for chan in eligible_channels:
|
|
name_counts[chan.name] = name_counts.get(chan.name, 0) + 1
|
|
|
|
for i, chan in enumerate(eligible_channels):
|
|
display_name = chan.name
|
|
if name_counts[chan.name] > 1:
|
|
cat_name = cat_map.get(chan.category_id)
|
|
if cat_name:
|
|
display_name = f"{chan.name} [{cat_name}]"
|
|
|
|
# Check for existing backup
|
|
found_prefix = ""
|
|
backup_file = self.exporter.export_path / "message_backup" / f"{chan.id}.json"
|
|
if backup_file.exists():
|
|
found_prefix = "[green][FOUND][/green] "
|
|
|
|
console.print(f"({i+1}) {found_prefix}{display_name}")
|
|
|
|
console.print("(A) [bold green]All Channels[/bold green]")
|
|
console.print("(B) Back")
|
|
|
|
choices = [str(i+1) for i in range(len(eligible_channels))] + ["A", "B", "a", "b"]
|
|
selection_input = Prompt.ask("\nSelect option or indices (e.g. 1,2,5)", default="B")
|
|
|
|
if selection_input.upper() == "B":
|
|
return
|
|
|
|
if selection_input.upper() == "A":
|
|
selected_channels = eligible_channels
|
|
else:
|
|
try:
|
|
# Handle multiple indices like '1,2,5'
|
|
indices = [int(i.strip()) - 1 for i in selection_input.split(",") if i.strip().isdigit()]
|
|
selected_channels = [eligible_channels[i] for i in indices if 0 <= i < len(eligible_channels)]
|
|
except Exception:
|
|
console.print("[red]Invalid selection. Aborting.[/red]")
|
|
return
|
|
|
|
if not selected_channels:
|
|
console.print("[yellow]No valid channels selected.[/yellow]")
|
|
return
|
|
|
|
# Check if any have [FOUND]
|
|
any_found = False
|
|
for chan in selected_channels:
|
|
if (self.exporter.export_path / "message_backup" / f"{chan.id}.json").exists():
|
|
any_found = True
|
|
break
|
|
|
|
if force_overwrite is None:
|
|
force_overwrite = False
|
|
if any_found and not skip_found_prompt:
|
|
console.print("\n[bold yellow]Existing backup(s) found for some selected channels.[/bold yellow]")
|
|
console.print("(Y) Update & Sync Backup")
|
|
console.print("(F) [bold red]Force Overwrite Backup[/bold red]")
|
|
console.print("(B) Back")
|
|
|
|
sync_choice = Prompt.ask("\nSelect option", choices=["Y", "F", "B"], default="Y").upper()
|
|
if sync_choice == "B":
|
|
return
|
|
force_overwrite = (sync_choice == "F")
|
|
|
|
console.print(f"\n[yellow]Starting backup for {len(selected_channels)} channels...[/yellow]")
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
TextColumn("- [bold yellow]{task.fields[msg_count]} {task.fields[suffix]}"),
|
|
console=console
|
|
) as progress:
|
|
for chan in selected_channels:
|
|
# Determine if it's a sync or fresh backup
|
|
backup_exists = (self.exporter.export_path / "message_backup" / f"{chan.id}.json").exists()
|
|
is_sync = backup_exists and not force_overwrite
|
|
|
|
label = "Syncing Backup" if is_sync else "Backing up"
|
|
suffix = "new messages" if is_sync else "messages"
|
|
|
|
# Create a specific task for each channel to keep it on its own line
|
|
task_id = progress.add_task(f"[cyan]{label}: {chan.name}", total=None, msg_count=0, suffix=suffix)
|
|
|
|
async def update_msg_count(name, count, tid=task_id):
|
|
progress.update(tid, msg_count=count)
|
|
|
|
await self.exporter.export_channel_messages(chan.id, progress_callback=update_msg_count, force=force_overwrite)
|
|
await self.exporter.export_threads(chan.id, progress_callback=update_msg_count, force=force_overwrite)
|
|
|
|
# Mark as finished (stop spinner)
|
|
progress.stop_task(task_id)
|
|
progress.update(task_id, description=f"[green]Completed: {chan.name}")
|
|
|
|
console.print("[bold green]Message backup complete![/bold green]")
|
|
# Update last_backup in server_profile.json
|
|
await self.exporter.export_metadata()
|
|
except Exception as e:
|
|
console.print(f"[bold red]Message backup failed: {e}[/bold red]")
|
|
finally:
|
|
await self.engine.close_connections()
|
|
|
|
async def sync_backup(self):
|
|
"""Option 3: Update & Sync Backup (Automated incremental sync for existing backups)."""
|
|
console.print("[yellow]Syncing backup... (Updating existing data)[/yellow]")
|
|
await self.backup_server_profile()
|
|
# Automatically select and sync existing backups
|
|
await self.backup_messages(auto_sync_existing=True)
|
|
|
|
async def edit_configuration(self):
|
|
# reuse or implement simplified version of edit_configuration from app.py
|
|
console.print("\n[bold]Configuration Editor[/bold]")
|
|
d_token = Prompt.ask("Discord Bot Token", default=self.config.discord_bot_token)
|
|
d_server = Prompt.ask("Discord Server ID", default=self.config.discord_server_id)
|
|
|
|
if d_token != self.config.discord_bot_token or d_server != self.config.discord_server_id:
|
|
self.config.discord_bot_token = d_token
|
|
self.config.discord_server_id = d_server
|
|
save_config(self.config, self.config_path)
|
|
self.engine = MigrationContext(self.config, target_platform="fluxer")
|
|
self.exporter = DiscordExporter(self.engine.discord_reader)
|
|
await self.validate_config()
|
|
console.print("[bold green]Config updated![/bold green]")
|
|
else:
|
|
console.print("[yellow]No changes.[/yellow]")
|
|
|
|
async def run_disco_reaper(config_path="config.yaml"):
|
|
app = DiscoReaperCLI(config_path)
|
|
await app.run()
|