discord_tools/scripts/backup_server.py

# discord_tools/scripts/server_backup.py

import os
import sys
import json
import asyncio
import aiohttp
from datetime import datetime
import requests

# Add the parent directory to the Python path
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(script_dir))
sys.path.insert(0, project_root)

from discord_tools.config.settings import DISCORD_TOKEN, DATA_DIR
from discord_tools.utils.api_utils import make_discord_request

class DiscordServerBackup:
    def __init__(self, guild_id):
        self.guild_id = guild_id
        self.backup_data = {}
        self.backup_folder = os.path.join(DATA_DIR, f"server_backup_{guild_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
        self.create_folder_structure()

    def create_folder_structure(self):
        folders = [
            self.backup_folder,
            os.path.join(self.backup_folder, "channels"),
            os.path.join(self.backup_folder, "images"),
            os.path.join(self.backup_folder, "images", "emojis"),
            os.path.join(self.backup_folder, "images", "stickers"),
        ]
        for folder in folders:
            os.makedirs(folder, exist_ok=True)

    def download_file(self, url, filepath):
        response = requests.get(url)
        if response.status_code == 200:
            with open(filepath, 'wb') as f:
                f.write(response.content)
            print(f"Downloaded {os.path.basename(filepath)}")
        else:
            print(f"Failed to download {os.path.basename(filepath)}")

    async def fetch_guild_info(self):
        endpoint = f'/guilds/{self.guild_id}'
        response = make_discord_request('GET', endpoint)
        if response:
            guild_info = response.json()
            self.save_json("guild_info", guild_info)
            print("Guild information fetched successfully.")
            
            # Download icon
            if 'icon' in guild_info:
                icon_url = f"https://cdn.discordapp.com/icons/{self.guild_id}/{guild_info['icon']}.png"
                self.download_file(icon_url, os.path.join(self.backup_folder, "images", "guild_icon.png"))
            
            # Download banner
            if 'banner' in guild_info:
                banner_url = f"https://cdn.discordapp.com/banners/{self.guild_id}/{guild_info['banner']}.png"
                self.download_file(banner_url, os.path.join(self.backup_folder, "images", "guild_banner.png"))
        else:
            print("Failed to fetch guild information.")

    async def fetch_channels(self):
        endpoint = f'/guilds/{self.guild_id}/channels'
        response = make_discord_request('GET', endpoint)
        if response:
            channels = response.json()
            self.save_json("channels", channels)
            
            # Create channel structure
            channel_structure = self.create_channel_structure(channels)
            self.save_json("channel_structure", channel_structure)
            
            # Fetch channel-specific settings
            await self.fetch_channel_settings(channels)
            
            print("Channel information and structure fetched successfully.")
        else:
            print("Failed to fetch channel information.")

    def create_channel_structure(self, channels):
        structure = {}
        categories = {c['id']: c for c in channels if c['type'] == 4}
        
        for channel in channels:
            if channel['type'] != 4:  # Not a category
                parent_id = channel.get('parent_id')
                if parent_id:
                    if parent_id not in structure:
                        structure[parent_id] = []
                    structure[parent_id].append(channel)
                else:
                    if 'no_category' not in structure:
                        structure['no_category'] = []
                    structure['no_category'].append(channel)
        
        return structure

    async def fetch_channel_settings(self, channels):
        for channel in channels:
            endpoint = f'/channels/{channel["id"]}'
            response = make_discord_request('GET', endpoint)
            if response:
                channel_data = response.json()
                self.save_json(f"channels/{channel['id']}_settings", channel_data)
            else:
                print(f"Failed to fetch settings for channel {channel['name']}")

    async def fetch_roles(self):
        endpoint = f'/guilds/{self.guild_id}/roles'
        response = make_discord_request('GET', endpoint)
        if response:
            roles = response.json()
            self.save_json("roles", roles)
            print("Role information fetched successfully.")
        else:
            print("Failed to fetch role information.")

    async def fetch_emojis(self):
        endpoint = f'/guilds/{self.guild_id}/emojis'
        response = make_discord_request('GET', endpoint)
        if response:
            emojis = response.json()
            self.save_json("emojis", emojis)
            print("Emoji information fetched successfully.")
            
            # Download custom emojis
            for emoji in emojis:
                emoji_url = f"https://cdn.discordapp.com/emojis/{emoji['id']}.png"
                self.download_file(emoji_url, os.path.join(self.backup_folder, "images", "emojis", f"{emoji['name']}.png"))
        else:
            print("Failed to fetch emoji information.")

    async def fetch_stickers(self):
        endpoint = f'/guilds/{self.guild_id}/stickers'
        response = make_discord_request('GET', endpoint)
        if response:
            stickers = response.json()
            self.save_json("stickers", stickers)
            print("Sticker information fetched successfully.")
            
            # Download custom stickers
            for sticker in stickers:
                sticker_url = f"https://cdn.discordapp.com/stickers/{sticker['id']}.png"
                self.download_file(sticker_url, os.path.join(self.backup_folder, "images", "stickers", f"{sticker['name']}.png"))
        else:
            print("Failed to fetch sticker information.")

    async def fetch_invites(self):
        endpoint = f'/guilds/{self.guild_id}/invites'
        response = make_discord_request('GET', endpoint)
        if response:
            invites = response.json()
            self.save_json("invites", invites)
            print("Server invites fetched successfully.")
        else:
            print("Failed to fetch server invites.")

    async def fetch_messages(self, channel_id, limit=None):
        messages = []
        last_message_id = None
        while True:
            endpoint = f'/channels/{channel_id}/messages?limit=100'
            if last_message_id:
                endpoint += f'&before={last_message_id}'
            response = make_discord_request('GET', endpoint)
            if response and response.status_code == 200:
                new_messages = response.json()
                if not new_messages:
                    break
                messages.extend(new_messages)
                last_message_id = new_messages[-1]['id']
                if limit and len(messages) >= limit:
                    messages = messages[:limit]
                    break
            else:
                print(f"Failed to fetch messages for channel {channel_id}")
                break
        return messages

    async def backup_messages(self):
        channels_data = self.load_json("channels")
        if not channels_data:
            print("Channel information not available. Skipping message backup.")
            return

        text_channels = [channel for channel in channels_data if channel['type'] == 0]

        for channel in text_channels:
            print(f"Backing up messages from channel: {channel['name']}")
            messages = await self.fetch_messages(channel['id'])
            self.save_json(f"channels/{channel['id']}_messages", messages)
            print(f"Backed up {len(messages)} messages from {channel['name']}")

    async def create_backup(self, include_messages=False):
        tasks = [
            self.fetch_guild_info(),
            self.fetch_channels(),
            self.fetch_roles(),
            self.fetch_emojis(),
            self.fetch_stickers(),
            self.fetch_invites(),
        ]
        await asyncio.gather(*tasks)

        if include_messages:
            await self.backup_messages()

    def save_json(self, filename, data):
        filepath = os.path.join(self.backup_folder, f"{filename}.json")
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=4)

    def load_json(self, filename):
        filepath = os.path.join(self.backup_folder, f"{filename}.json")
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            print(f"File not found: {filepath}")
            return None

async def main():
    guild_id = input("Enter the Discord server ID to backup: ").strip()
    
    if not guild_id.isdigit():
        print("Invalid server ID. Please enter a numeric ID.")
        return

    include_messages = input("Do you want to backup all messages? (y/n): ").strip().lower() == 'y'

    backup = DiscordServerBackup(guild_id)
    await backup.create_backup(include_messages)
    print(f"Backup completed. Files saved in {backup.backup_folder}")

if __name__ == "__main__":
    asyncio.run(main())