get_server_icons.py

import sys
import json
import logging
import secrets
import requests
import webbrowser
import urllib.parse
import http.server
import socketserver
from pathlib import Path

from config import load_environment, validate_config, ensure_env_file, ROOT

# Configure logging
logger = logging.getLogger(__name__)

# Validate environment
if not ensure_env_file():
    sys.exit(1)

config = load_environment()
if not validate_config(config, ['oauth_client_id', 'oauth_client_secret']):
    sys.exit(1)

CLIENT_ID = config['oauth_client_id']
CLIENT_SECRET = config['oauth_client_secret']
REDIRECT_URI = "http://localhost:8080/callback"
SCOPES = "identify guilds"

def main():
    """Fetch server icon URLs using OAuth2 flow with CSRF protection."""
    logger.info("Starting Discord OAuth2 flow for guild discovery...")

    # Generate state token for CSRF protection
    state_token = secrets.token_urlsafe(32)
    logger.debug(f"Generated state token: {state_token}")

    token_url = "https://discord.com/api/oauth2/token"
    guilds_url = "https://discord.com/api/v10/users/@me/guilds"

    # Build authorization URL with state parameter
    params = {
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "response_type": "code",
        "scope": SCOPES,
        "state": state_token,
    }
    auth_url = f"https://discord.com/api/oauth2/authorize?{urllib.parse.urlencode(params)}"
    logger.info(f"Opening authorization URL: {auth_url}")
    webbrowser.open(auth_url)

    # OAuth callback handler with state validation
    class OAuthHandler(http.server.SimpleHTTPRequestHandler):
        def log_message(self, format, *args):
            """Override logging to use logger instead of print."""
            logger.debug(format % args)

        def do_GET(self):
            """Handle OAuth redirect callback."""
            query = urllib.parse.urlparse(self.path).query
            params = urllib.parse.parse_qs(query)

            # Validate state parameter for CSRF protection
            returned_state = params.get("state", [None])[0]
            if returned_state != state_token:
                logger.error(f"State mismatch! Expected {state_token}, got {returned_state}")
                self.send_response(400)
                self.end_headers()
                self.wfile.write(b"Invalid state parameter - possible CSRF attack")
                return

            self.server.auth_code = params.get("code", [None])[0]

            if not self.server.auth_code:
                logger.error("No authorization code received")
                self.send_response(400)
                self.end_headers()
                self.wfile.write(b"No authorization code received")
                return

            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"Authorization successful! You can close this tab.")
            logger.info("Authorization successful")

    # Start local server to receive OAuth callback
    logger.info("Starting local callback server on localhost:8080...")
    try:
        with socketserver.TCPServer(("localhost", 8080), OAuthHandler) as httpd:
            httpd.handle_request()
            auth_code = httpd.auth_code
    except Exception as e:
        logger.error(f"Failed to start callback server: {e}")
        sys.exit(1)

    if not auth_code:
        logger.error("No authorization code received from Discord")
        sys.exit(1)

    # Exchange authorization code for access token
    logger.info("Exchanging authorization code for access token...")
    token_data = {
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        "grant_type": "authorization_code",
        "code": auth_code,
        "redirect_uri": REDIRECT_URI,
    }
    token_headers = {"Content-Type": "application/x-www-form-urlencoded"}

    try:
        token_resp = requests.post(token_url, data=token_data, headers=token_headers, timeout=10)
        token_resp.raise_for_status()
    except requests.RequestException as e:
        logger.error(f"Failed to exchange code for token: {e}")
        sys.exit(1)

    try:
        access_token = token_resp.json()["access_token"]
    except (KeyError, json.JSONDecodeError) as e:
        logger.error(f"Invalid token response from Discord: {e}")
        sys.exit(1)

    logger.info("Access token obtained successfully")

    # Fetch guilds with the access token
    logger.info("Fetching guild list...")
    headers = {"Authorization": f"Bearer {access_token}"}

    try:
        guilds_resp = requests.get(guilds_url, headers=headers, timeout=10)
        guilds_resp.raise_for_status()
        guilds = guilds_resp.json()
    except requests.RequestException as e:
        logger.error(f"Failed to fetch guilds: {e}")
        sys.exit(1)

    logger.info(f"Found {len(guilds)} guilds")

    # Extract server icons
    result = {}
    for guild in guilds:
        icon = guild.get("icon")
        if not icon:
            logger.debug(f"Guild {guild.get('name', guild.get('id'))} has no icon")
            continue

        ext = "gif" if icon.startswith("a_") else "png"
        icon_url = f"https://cdn.discordapp.com/icons/{guild['id']}/{icon}.{ext}?size=512"
        result[guild["id"]] = icon_url
        logger.debug(f"Guild {guild.get('name', guild['id'])}: {icon_url}")

    # Save results to servers.json
    out_file = ROOT / 'servers.json'
    try:
        with open(out_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2)
        logger.info(f"Saved {len(result)} guild icon URLs to {out_file}")
    except Exception as e:
        logger.error(f"Failed to save servers.json: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()