ctfdown.py

#!/usr/bin/env python3
"""
CTFdown - CTFd Challenge Downloader
A comprehensive tool for downloading and managing CTF challenges from CTFd platforms.
"""

import argparse
import hashlib
import json
import os
import re
import sys
import time
import zipfile
import tarfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from html import unescape
from pathlib import Path
from urllib.parse import urljoin, urlparse

try:
    from tqdm import tqdm
    HAS_TQDM = True
except ImportError:
    HAS_TQDM = False

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry


class CTFdown:
    def __init__(self, base_url, session_cookie, max_workers=4, max_retries=3, retry_backoff=1, output_dir=None):
        self.base_url = base_url.rstrip('/')
        self.session = self._create_session(session_cookie, max_retries, retry_backoff)
        self.max_workers = max_workers
        self.output_dir = output_dir or 'challenges'
        self.state_file = Path(self.output_dir) / '.ctfdown_state.json'
        self.state = self._load_state()
    
    def _create_session(self, session_cookie, max_retries, retry_backoff):
        session = requests.Session()
        session.cookies.set('session', session_cookie)
        session.headers.update({
            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36'
        })
        
        retry_strategy = Retry(
            total=max_retries,
            backoff_factor=retry_backoff,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET", "POST"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        
        return session
    
    def _load_state(self):
        if self.state_file.exists():
            try:
                with open(self.state_file, 'r') as f:
                    return json.load(f)
            except:
                return {}
        return {}
    
    def _save_state(self):
        try:
            with open(self.state_file, 'w') as f:
                json.dump(self.state, f, indent=2)
        except:
            pass
    
    def _get_file_hash(self, filepath):
        if not Path(filepath).exists():
            return None
        sha256 = hashlib.sha256()
        with open(filepath, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b''):
                sha256.update(chunk)
        return sha256.hexdigest()
    
    def test_connection(self):
        try:
            response = self.session.get(f"{self.base_url}/api/v1/challenges", timeout=10)
            if response.status_code == 200:
                return True
            return False
        except:
            return False
    
    def get_challenges(self, filters=None):
        try:
            response = self.session.get(f"{self.base_url}/api/v1/challenges", timeout=30)
            if response.status_code != 200:
                return []
            
            data = response.json()
            challenges = data.get('data', []) if isinstance(data, dict) else data
            
            if filters:
                challenges = self._filter_challenges(challenges, filters)
            
            return challenges
        except Exception as e:
            print(f"✗ Error fetching challenges: {e}")
            return []
    
    def _filter_challenges(self, challenges, filters):
        filtered = challenges
        
        if filters.get('category'):
            categories = [c.strip() for c in filters['category'].split(',')]
            filtered = [c for c in filtered if c.get('category', '').lower() in [cat.lower() for cat in categories]]
        
        if filters.get('tags'):
            tags = [t.strip() for t in filters['tags'].split(',')]
            filtered = [c for c in filtered if any(
                tag.lower() in str(t.get('value', t) if isinstance(t, dict) else t).lower()
                for t in c.get('tags', []) for tag in tags
            )]
        
        if filters.get('solved_only'):
            filtered = [c for c in filtered if c.get('solved_by_me', False)]
        
        if filters.get('unsolved_only'):
            filtered = [c for c in filtered if not c.get('solved_by_me', False)]
        
        if filters.get('min_points'):
            filtered = [c for c in filtered if c.get('value', 0) >= filters['min_points']]
        
        if filters.get('max_points'):
            filtered = [c for c in filtered if c.get('value', 0) <= filters['max_points']]
        
        if filters.get('challenge_ids'):
            ids = [int(i.strip()) for i in filters['challenge_ids'].split(',')]
            filtered = [c for c in filtered if c.get('id') in ids]
        
        return filtered
    
    def get_challenge_details(self, challenge_id):
        try:
            response = self.session.get(f"{self.base_url}/api/v1/challenges/{challenge_id}", timeout=30)
            if response.status_code == 200:
                data = response.json()
                return data.get('data') if isinstance(data, dict) else data
            return None
        except:
            return None
    
    def _should_update_challenge(self, challenge_path, challenge_id, challenge_data):
        if not challenge_path.exists():
            return True
        
        state_key = str(challenge_id)
        stored_hash = self.state.get(state_key, {}).get('hash')
        
        current_hash = hashlib.sha256(
            json.dumps(challenge_data, sort_keys=True).encode()
        ).hexdigest()
        
        if stored_hash != current_hash:
            return True
        
        title_file = challenge_path / "Title_Description.txt"
        if not title_file.exists():
            return True
        
        return False
    
    def _extract_description(self, view_html, description):
        if view_html:
            desc_match = re.search(
                r'<span[^>]*class=["\']challenge-desc["\'][^>]*>(.*?)</span>',
                view_html, re.DOTALL | re.IGNORECASE
            )
            if desc_match:
                description = desc_match.group(1)
        return description
    
    def _extract_connection_info(self, view_html, connection_info):
        if not connection_info and view_html:
            conn_match = re.search(
                r'<span[^>]*class=["\']challenge-connection-info["\'][^>]*>(.*?)</span>',
                view_html, re.DOTALL | re.IGNORECASE
            )
            if conn_match:
                conn_html = conn_match.group(1)
                conn_text = re.sub(r'<[^>]+>', '', conn_html).strip()
                if conn_text:
                    connection_info = conn_text
        return connection_info
    
    def _process_html_description(self, description):
        if not description or ('<' not in description and '&' not in description and 'href=' not in description.lower()):
            return description
        
        def replace_link(match):
            url = match.group(1)
            link_text = match.group(2) if len(match.groups()) >= 2 else ''
            if not link_text or not link_text.strip():
                link_text = url
            return url if link_text.strip() == url else f"{link_text.strip()} ({url})"
        
        description = re.sub(
            r'<a\s+[^>]*href=["\']([^"\']+)["\'][^>]*>([^<]*)</a>',
            replace_link, description, flags=re.IGNORECASE
        )
        description = re.sub(
            r'<a\s+[^>]*href=["\']([^"\']+)["\'][^>]*>',
            lambda m: f"{m.group(1)} ", description, flags=re.IGNORECASE
        )
        description = re.sub(r'<[^>]+>', '', description)
        description = unescape(description)
        description = re.sub(r'\s+', ' ', description).strip()
        description = description.replace(' <br> ', '\n').replace('<br>', '\n')
        description = description.replace(' <p> ', '\n\n').replace('</p>', '\n\n')
        description = re.sub(r'\n{3,}', '\n\n', description)
        
        return description
    
    def _write_metadata_file(self, challenge_path, challenge):
        title = challenge.get('name', 'Unknown Challenge')
        view_html = challenge.get('view', '')
        description = self._extract_description(view_html, challenge.get('description', 'No description available.'))
        connection_info = self._extract_connection_info(view_html, challenge.get('connection_info', ''))
        description = self._process_html_description(description)
        
        title_desc_file = challenge_path / "Title_Description.txt"
        with open(title_desc_file, 'w', encoding='utf-8') as f:
            f.write(f"Title: {title}\n")
            
            for key, label in [('category', 'Category'), ('type', 'Type'), ('value', 'Points')]:
                value = challenge.get(key, '')
                if value:
                    f.write(f"{label}: {value}\n")
            
            solves = challenge.get('solves', 0)
            if solves is not None:
                f.write(f"Solves: {solves}\n")
            
            tags = challenge.get('tags', [])
            if tags:
                tag_values = [t.get('value', str(t)) if isinstance(t, dict) else str(t) for t in tags]
                if tag_values:
                    f.write(f"Tags: {', '.join(tag_values)}\n")
            
            attribution = challenge.get('attribution', '').strip()
            if attribution:
                f.write(f"Attribution: {attribution}\n")
            
            f.write(f"\nDescription:\n{description}\n")
            
            if connection_info:
                f.write(f"\nConnection Info:\n{connection_info}\n")
            
            hints = challenge.get('hints', [])
            if hints:
                f.write(f"\nHints:\n")
                for i, hint in enumerate(hints, 1):
                    hint_content = hint.get('content', hint.get('hint', str(hint))) if isinstance(hint, dict) else hint
                    f.write(f"  {i}. {hint_content}\n")
            
            solution_state = challenge.get('solution_state', 'hidden')
            if solution_state == 'visible' and view_html:
                solution_match = re.search(
                    r'<div[^>]*class=["\']challenge-solution-content["\'][^>]*>(.*?)</div>',
                    view_html, re.DOTALL | re.IGNORECASE
                )
                if solution_match:
                    solution_html = solution_match.group(1)
                    solution_text = re.sub(r'<[^>]+>', '', solution_html)
                    solution_text = unescape(solution_text).strip()
                    if solution_text:
                        f.write(f"\nSolution:\n{solution_text}\n")
    
    def _extract_filename(self, file_url, idx):
        parsed = urlparse(file_url)
        path_parts = [p for p in parsed.path.strip('/').split('/') if p]
        
        if len(path_parts) >= 2:
            filename = path_parts[-1]
        elif len(path_parts) == 1:
            filename = path_parts[0]
        else:
            filename = f"file_{idx}"
        
        filename = filename.split('?')[0]
        
        if not filename or (len(filename) == 32 and all(c in '0123456789abcdef' for c in filename.lower())):
            filename = f"file_{idx}"
        
        return filename
    
    def download_file(self, url, filepath, progress_bar=None):
        try:
            response = self.session.get(url, stream=True, timeout=60)
            if response.status_code == 200:
                content_disposition = response.headers.get('Content-Disposition', '')
                if 'filename=' in content_disposition:
                    match = re.search(r'filename="?([^"]+)"?', content_disposition)
                    if match:
                        filepath = os.path.join(os.path.dirname(filepath), match.group(1))
                
                os.makedirs(os.path.dirname(filepath), exist_ok=True)
                total_size = int(response.headers.get('content-length', 0))
                
                with open(filepath, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            if progress_bar and total_size > 0:
                                try:
                                    progress_bar.update(len(chunk))
                                except:
                                    pass
                
                return self._get_file_hash(filepath)
            return None
        except Exception as e:
            if progress_bar:
                try:
                    progress_bar.write(f"✗ Error downloading {url}: {e}")
                except:
                    pass
            return None
    
    def _validate_challenge(self, challenge_path):
        errors = []
        warnings = []
        
        if not challenge_path.exists():
            errors.append("Challenge directory does not exist")
            return errors, warnings
        
        title_file = challenge_path / "Title_Description.txt"
        if not title_file.exists():
            errors.append("Title_Description.txt missing")
        else:
            with open(title_file, 'r') as f:
                content = f.read()
                if 'Title:' not in content:
                    errors.append("Title missing in Title_Description.txt")
                if 'Description:' not in content:
                    warnings.append("Description missing in Title_Description.txt")
        
        files_dir = challenge_path / "challenge_files"
        if files_dir.exists():
            files = list(files_dir.iterdir())
            if not files:
                warnings.append("challenge_files directory is empty")
            else:
                for file in files:
                    if not file.is_file():
                        warnings.append(f"Non-file item in challenge_files: {file.name}")
                    elif file.stat().st_size == 0:
                        warnings.append(f"Empty file: {file.name}")
        
        return errors, warnings
    
    def sanitize_filename(self, name):
        invalid_chars = '<>:"/\\|?*'
        for char in invalid_chars:
            name = name.replace(char, '_')
        return name.strip(' .')
    
    def download_challenges(self, output_dir="challenges", filters=None, update_mode=False, parallel=False, validate=False, progress=True):
        challenges = self.get_challenges(filters)
        if not challenges:
            print("No challenges found.")
            return None
        
        base_path = Path(output_dir)
        base_path.mkdir(exist_ok=True)
        
        categories = {}
        for challenge in challenges:
            category = challenge.get('category', 'Uncategorized') or 'Uncategorized'
            categories.setdefault(category, []).append(challenge)
        
        total_files = 0
        updated = 0
        skipped = 0
        failed = 0
        
        progress_bar = None
        if progress and HAS_TQDM:
            progress_bar = tqdm(total=len(challenges), desc="Challenges", unit="challenge")
        
        try:
            for category_name, category_challenges in categories.items():
                sanitized_category = self.sanitize_filename(category_name)
                category_path = base_path / sanitized_category
                category_path.mkdir(exist_ok=True)
                
                for challenge in category_challenges:
                    try:
                        challenge_id = challenge.get('id')
                        challenge_name = challenge.get('name', f"Challenge_{challenge_id}")
                        sanitized_challenge = self.sanitize_filename(challenge_name)
                        challenge_path = category_path / sanitized_challenge
                        challenge_path.mkdir(exist_ok=True)
                        
                        full_details = self.get_challenge_details(challenge_id) if challenge_id else challenge
                        if not full_details:
                            full_details = challenge
                        
                        if update_mode and not self._should_update_challenge(challenge_path, challenge_id, full_details):
                            skipped += 1
                            if progress_bar:
                                progress_bar.update(1)
                            continue
                        
                        self._write_metadata_file(challenge_path, full_details)
                        
                        files = full_details.get('files', [])
                        if files:
                            files_path = challenge_path / "challenge_files"
                            files_path.mkdir(exist_ok=True)
                            
                            file_tasks = []
                            for idx, file_url in enumerate(files, 1):
                                if not file_url.startswith('http'):
                                    full_url = urljoin(self.base_url, file_url)
                                else:
                                    full_url = file_url
                                
                                filename = self._extract_filename(full_url, idx)
                                file_path = files_path / filename
                                
                                file_tasks.append((full_url, str(file_path), filename))
                            
                            if parallel and len(file_tasks) > 1:
                                file_progress = None
                                if progress and HAS_TQDM:
                                    try:
                                        file_progress = tqdm(total=len(file_tasks), desc=f"  {challenge_name[:30]}", leave=False, unit="file")
                                    except:
                                        file_progress = None
                                
                                with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                                    futures = {executor.submit(self.download_file, url, path, None): (url, path, name) 
                                             for url, path, name in file_tasks}
                                    
                                    for future in as_completed(futures):
                                        url, path, name = futures[future]
                                        try:
                                            file_hash = future.result()
                                            if file_hash:
                                                total_files += 1
                                                state_key = f"{challenge_id}_file_{name}"
                                                self.state[state_key] = {'hash': file_hash}
                                                if file_progress:
                                                    try:
                                                        file_progress.update(1)
                                                    except:
                                                        pass
                                            else:
                                                failed += 1
                                                if file_progress:
                                                    try:
                                                        file_progress.update(1)
                                                    except:
                                                        pass
                                        except Exception as e:
                                            failed += 1
                                            if file_progress:
                                                try:
                                                    file_progress.write(f"      ✗ Error: {e}")
                                                    file_progress.update(1)
                                                except:
                                                    pass
                                
                                if file_progress:
                                    try:
                                        file_progress.close()
                                    except:
                                        pass
                            else:
                                for url, path, name in file_tasks:
                                    file_progress = None
                                    if progress and HAS_TQDM:
                                        try:
                                            file_progress = tqdm(desc=f"    {name[:40]}", unit="B", unit_scale=True, leave=False, total=None)
                                        except:
                                            file_progress = None
                                    
                                    file_hash = self.download_file(url, path, file_progress)
                                    if file_hash:
                                        total_files += 1
                                        state_key = f"{challenge_id}_file_{name}"
                                        self.state[state_key] = {'hash': file_hash}
                                    else:
                                        failed += 1
                                    
                                    if file_progress:
                                        try:
                                            file_progress.close()
                                        except:
                                            pass
                        
                        challenge_hash = hashlib.sha256(
                            json.dumps(full_details, sort_keys=True).encode()
                        ).hexdigest()
                        self.state[str(challenge_id)] = {'hash': challenge_hash}
                        updated += 1
                        
                        if validate:
                            errors, warnings = self._validate_challenge(challenge_path)
                            if errors:
                                print(f"    ✗ Validation errors for {challenge_name}: {', '.join(errors)}")
                            if warnings:
                                print(f"    ⚠ Validation warnings for {challenge_name}: {', '.join(warnings)}")
                        
                        if progress_bar:
                            progress_bar.update(1)
                    
                    except Exception as e:
                        failed += 1
                        print(f"    ✗ Error processing {challenge_name}: {e}")
                        if progress_bar:
                            progress_bar.update(1)
            
            self._save_state()
            
            if progress_bar:
                progress_bar.close()
            
            print(f"\n=== Download Complete ===")
            print(f"Updated: {updated}, Skipped: {skipped}, Failed: {failed}")
            print(f"Total files downloaded: {total_files}")
            
            return base_path
        
        except KeyboardInterrupt:
            if progress_bar:
                progress_bar.close()
            print("\n\n⚠ Download interrupted. Partial download saved.")
            self._save_state()
            return base_path
    
    def create_archive(self, source_dir, archive_path, archive_format='zip'):
        source_path = Path(source_dir)
        archive_path = Path(archive_path)
        
        if not source_path.exists():
            print(f"✗ Source directory does not exist: {source_dir}")
            return False
        
        try:
            if archive_format == 'zip':
                with zipfile.ZipFile(archive_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                    files = []
                    for root, dirs, filenames in os.walk(source_path):
                        dirs[:] = [d for d in dirs if not d.startswith('.')]
                        for filename in [f for f in filenames if not f.startswith('.')]:
                            file_path = Path(root) / filename
                            arcname = file_path.relative_to(source_path.parent)
                            zipf.write(file_path, arcname)
                            files.append(arcname)
                    print(f"✓ Created ZIP archive: {archive_path} ({len(files)} files)")
            
            elif archive_format in ('tar.gz', 'tgz'):
                with tarfile.open(archive_path, 'w:gz') as tar:
                    def filter_func(tarinfo):
                        return None if tarinfo.name.startswith('.') else tarinfo
                    
                    tar.add(source_path, arcname=source_path.name, recursive=True, filter=filter_func)
                    file_count = sum(1 for m in tar.getmembers() if m.isfile())
                    print(f"✓ Created TAR.GZ archive: {archive_path} ({file_count} files)")
            
            archive_size = archive_path.stat().st_size
            size_mb = archive_size / (1024 * 1024)
            print(f"  Size: {size_mb:.2f} MB ({archive_size:,} bytes)")
            return True
        
        except Exception as e:
            print(f"✗ Error creating archive: {e}")
            if archive_path.exists():
                archive_path.unlink()
            return False


def main():
    parser = argparse.ArgumentParser(
        prog='ctfdown',
        description='CTFdown - Download and manage CTF challenges from CTFd platforms',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  ctfdown https://ctf.example.com SESSION_COOKIE
  ctfdown https://ctf.example.com SESSION_COOKIE --zip --archive-name MyCTF
  ctfdown https://ctf.example.com SESSION_COOKIE --category pwn,web --parallel
  ctfdown https://ctf.example.com SESSION_COOKIE --update --validate
  ctfdown https://ctf.example.com SESSION_COOKIE --unsolved-only --min-points 100
        """
    )
    
    parser.add_argument('url', help='CTFd platform base URL')
    parser.add_argument('session_cookie', help='Session cookie value for authentication')
    parser.add_argument('-o', '--output-dir', default='challenges', help='Output directory (default: challenges)')
    
    group_filter = parser.add_argument_group('filtering options')
    group_filter.add_argument('--category', help='Filter by category (comma-separated)')
    group_filter.add_argument('--tags', help='Filter by tags (comma-separated)')
    group_filter.add_argument('--solved-only', action='store_true', help='Only download solved challenges')
    group_filter.add_argument('--unsolved-only', action='store_true', help='Only download unsolved challenges')
    group_filter.add_argument('--min-points', type=int, help='Minimum points filter')
    group_filter.add_argument('--max-points', type=int, help='Maximum points filter')
    group_filter.add_argument('--challenge-ids', help='Specific challenge IDs (comma-separated)')
    
    group_mode = parser.add_argument_group('mode options')
    group_mode.add_argument('--update', action='store_true', help='Update mode: only download new/changed challenges')
    group_mode.add_argument('--parallel', action='store_true', help='Download files in parallel (faster but more resource-intensive)')
    group_mode.add_argument('--validate', action='store_true', help='Validate downloaded challenges')
    group_mode.add_argument('--inspect', action='store_true', help='Inspect API structure without downloading')
    group_mode.add_argument('--no-progress', action='store_true', help='Disable progress bars')
    
    group_archive = parser.add_argument_group('archive options')
    group_archive.add_argument('--zip', action='store_true', help='Create ZIP archive after downloading')
    group_archive.add_argument('--tar', '--tar.gz', '--tgz', dest='tar', action='store_true', help='Create TAR.GZ archive')
    group_archive.add_argument('--archive-name', default='CTF_Challenges', help='Archive name (default: CTF_Challenges)')
    
    group_advanced = parser.add_argument_group('advanced options')
    group_advanced.add_argument('--max-workers', type=int, default=4, help='Max parallel workers (default: 4)')
    group_advanced.add_argument('--max-retries', type=int, default=3, help='Max retry attempts (default: 3)')
    group_advanced.add_argument('--retry-backoff', type=float, default=1.0, help='Retry backoff factor (default: 1.0)')
    
    args = parser.parse_args()
    
    if not HAS_TQDM and not args.no_progress:
        print("⚠ tqdm not installed. Install with: pip install tqdm")
        print("   Continuing without progress bars...\n")
    
    downloader = CTFdown(
        args.url,
        args.session_cookie,
        max_workers=args.max_workers,
        max_retries=args.max_retries,
        retry_backoff=args.retry_backoff,
        output_dir=args.output_dir
    )
    
    if not downloader.test_connection():
        print("✗ Failed to connect. Please check your URL and session cookie.")
        sys.exit(1)
    
    if args.inspect:
        print("API inspection not implemented in this version.")
        sys.exit(0)
    
    filters = {}
    if args.category:
        filters['category'] = args.category
    if args.tags:
        filters['tags'] = args.tags
    if args.solved_only:
        filters['solved_only'] = True
    if args.unsolved_only:
        filters['unsolved_only'] = True
    if args.min_points:
        filters['min_points'] = args.min_points
    if args.max_points:
        filters['max_points'] = args.max_points
    if args.challenge_ids:
        filters['challenge_ids'] = args.challenge_ids
    
    challenges_path = downloader.download_challenges(
        output_dir=args.output_dir,
        filters=filters if filters else None,
        update_mode=args.update,
        parallel=args.parallel,
        validate=args.validate,
        progress=not args.no_progress
    )
    
    if (args.zip or args.tar) and challenges_path:
        archive_format = 'zip' if args.zip else 'tar.gz'
        archive_name = args.archive_name
        if not archive_name.endswith(('.zip', '.tar.gz', '.tgz')):
            archive_name += ('.zip' if args.zip else '.tar.gz')
        
        downloader.create_archive(challenges_path, archive_name, archive_format)


if __name__ == "__main__":
    main()