Rafael Marcos Luque Baena

"""
Database query functions for REDB.

Contains all functions that query ClickHouse for sample metadata,
deduplication checks, and catalog lookups.
"""
import os
import ast
import clickhouse_connect
from typing import List, Optional, Dict

from redb import settings
from redb.s3_utils import generate_s3_key_from_hash


def get_supported_formats(magika_filter: Optional[str] = None) -> List[str]:
    """
    Get supported file formats from SUPPORTED_FORMATS env variable or magika_filter override.
    Expected format: SUPPORTED_FORMATS=['pebin', 'elf']

    Args:
        magika_filter: Optional single format to filter by (overrides env var)

    Returns:
        List of supported format strings, defaults to ['pebin'] if not set
    """
    # If magika_filter is provided, use it as the only format
    if magika_filter:
        return [magika_filter]

    formats_str = os.getenv('SUPPORTED_FORMATS', "['pebin']")
    try:
        formats = ast.literal_eval(formats_str)
        if isinstance(formats, list) and all(isinstance(f, str) for f in formats):
            return formats
        else:
            print(f"[WARNING] SUPPORTED_FORMATS must be a list of strings, got: {formats_str}")
            return ['pebin']
    except (ValueError, SyntaxError) as e:
        print(f"[WARNING] Failed to parse SUPPORTED_FORMATS '{formats_str}': {e}")
        return ['pebin']


def get_db_catalog_connection():
    """Create and return a ClickHouse client for catalog queries"""
    return clickhouse_connect.get_client(
        host=os.getenv("DB_HOST"),
        port=int(os.getenv("DB_PORT", "8123")),
        verify=os.getenv("DB_ENFORCE_SSL", "False").lower() == "true",
        username=os.getenv("DB_USER"),
        password=os.getenv("DB_PASSWORD"),
        database=os.getenv("DB_NAME")
    )


def fetch_s3_objects_by_repository(repository: Optional[str], index_prefix: str, decompile: bool, notes: Optional[str] = None, magika_filter: Optional[str] = None, yara_scan: bool = False, force: bool = False) -> List[Dict]:
    """
    Query the ClickHouse repository_upload_sessions table to get samples.
    S3 bucket comes from env var, S3 key is derived from sha256.

    Args:
        repository: Repository name (e.g., "bazaar", "malshare"). If None or "all-repos", queries all repositories.
        index_prefix: Table prefix for checking existing samples
        decompile: Whether we're in decompile mode (affects which table to check for existing)
        notes: Optional filter for notes field
        magika_filter: Optional single filetype to filter by (overrides SUPPORTED_FORMATS)
        yara_scan: Whether we're in YARA-only mode (skips "already processed" check)

    Returns:
        List of dicts with sha256, s3_bucket, s3_key for each sample
    """
    try:
        client = get_db_catalog_connection()
        s3_bucket = os.getenv('S3_BUCKET')

        if not s3_bucket:
            print("[ERROR] S3_BUCKET environment variable is required")
            return []

        # Get supported formats from env or magika_filter override
        supported_formats = get_supported_formats(magika_filter)
        print(f"[INFO] Querying for supported formats: {supported_formats}")

        # Build query - repository filter is optional
        # Join with catalog_samples to get first_seen date
        if repository and repository != "all-repos":
            query = """
                SELECT DISTINCT rus.sha256, rus.filetype_magika, cs.first_seen
                FROM repository_upload_sessions rus
                LEFT JOIN catalog_samples cs ON rus.sha256 = cs.sha256
                WHERE rus.repository = %(repo)s
                  AND rus.filetype_magika IN %(formats)s
            """
            params = {"repo": repository, "formats": supported_formats}
        else:
            # Query all repositories
            query = """
                SELECT DISTINCT rus.sha256, rus.filetype_magika, cs.first_seen
                FROM repository_upload_sessions rus
                LEFT JOIN catalog_samples cs ON rus.sha256 = cs.sha256
                WHERE rus.filetype_magika IN %(formats)s
            """
            params = {"formats": supported_formats}

        if notes:
            query += " AND rus.notes LIKE %(notes)s"
            params["notes"] = f"%{notes}%"

        # Execute query and convert to list of dictionaries
        result = client.query(query, parameters=params)

        # Build rows with s3_bucket and s3_key derived from sha256
        rows = []
        for row in result.result_rows:
            sha256 = row[0]
            # Handle binary string if needed
            if isinstance(sha256, bytes):
                sha256 = sha256.decode('utf-8')

            rows.append({
                'sha256': sha256,
                's3_bucket': s3_bucket,
                's3_key': generate_s3_key_from_hash(sha256),
                'filetype_magika': row[1],
                'first_seen': row[2]  # From catalog_samples (None if not found)
            })

        print(f"[DEBUG] Found {len(rows)} objects in repository_upload_sessions for repository {repository}")

        # Extract SHA256 hashes from the results for bulk checking
        sha256_list = [row.get('sha256') for row in rows if row.get('sha256')]

        if sha256_list and not force:
            # Check which hashes are already in the database
            existing_hashes = is_in_db_bulk(sha256_list, index_prefix, decompile, yara_scan, magika_filter)
            print(f"[INFO] Found {len(existing_hashes)} objects already in database")

            # Filter out rows with existing hashes
            rows = [row for row in rows if row.get('sha256') not in existing_hashes]
            print(f"[INFO] After filtering, {len(rows)} objects remain to be processed")
        elif force:
            print(f"[INFO] Force mode: skipping deduplication check, processing all {len(sha256_list)} objects")

        # Randomize the order of rows before returning
        import random
        random.shuffle(rows)

        return rows
    except Exception as e:
        print(f"[ERROR] Failed to query repository_upload_sessions: {e}")
        return []
    finally:
        if 'client' in locals():
            client.close()


def fetch_s3_objects_by_date_range(
    index_prefix: str,
    decompile: bool,
    start_date: str,
    end_date: str,
    repository: Optional[str] = None,
    notes: Optional[str] = None,
    magika_filter: Optional[str] = None,
    yara_scan: bool = False,
    force: bool = False,
    analyzed: bool = False
) -> List[Dict]:
    """
    Query samples from catalog_samples by first_seen date range, filtered to only include
    samples that exist in repository_upload_sessions (bulk/repo uploads only).

    Args:
        index_prefix: Table prefix for checking existing samples
        decompile: Whether we're in decompile mode
        start_date: Start date (inclusive) in YYYY-MM-DD format
        end_date: End date (exclusive) in YYYY-MM-DD format
        repository: Optional repository filter
        notes: Optional notes filter
        magika_filter: Optional single filetype to filter by (overrides SUPPORTED_FORMATS)
        yara_scan: Whether we're in YARA-only mode (skips "already processed" check)
        analyzed: Filter to only samples already in basic_properties (cross-database join)

    Returns:
        List of dicts with sha256, s3_bucket, s3_key for each sample
    """
    try:
        client = get_db_catalog_connection()
        s3_bucket = os.getenv('S3_BUCKET')

        if not s3_bucket:
            print("[ERROR] S3_BUCKET environment variable is required")
            return []

        # Get supported formats from env or magika_filter override
        supported_formats = get_supported_formats(magika_filter)
        print(f"[INFO] Querying for supported formats: {supported_formats}")

        # Join catalog_samples with repository_upload_sessions to:
        # 1. Filter by first_seen date from catalog_samples
        # 2. Only include samples that exist in repository_upload_sessions (not user uploads)
        # 3. Optionally filter to only already-analyzed samples (in basic_properties)
        analyzed_join = ""
        if analyzed:
            basic_table = f"{index_prefix}_basic_properties"
            analyzed_join = f"INNER JOIN {basic_table} bp ON cs.sha256 = bp.sha256"
            print(f"[INFO] Filtering to already-analyzed samples in {basic_table}")

        query = f"""
            SELECT DISTINCT cs.sha256, rus.filetype_magika, cs.first_seen
            FROM catalog_samples cs
            INNER JOIN repository_upload_sessions rus ON cs.sha256 = rus.sha256
            {analyzed_join}
            WHERE cs.first_seen >= %(start_date)s
              AND cs.first_seen < %(end_date)s
              AND rus.filetype_magika IN %(formats)s
        """
        params = {
            "start_date": start_date,
            "end_date": end_date,
            "formats": supported_formats
        }

        if repository:
            query += " AND rus.repository = %(repo)s"
            params["repo"] = repository

        if notes:
            query += " AND rus.notes LIKE %(notes)s"
            params["notes"] = f"%{notes}%"

        # Execute query
        result = client.query(query, parameters=params)

        # Build rows with s3_bucket and s3_key derived from sha256
        rows = []
        for row in result.result_rows:
            sha256 = row[0]
            # Handle binary string if needed
            if isinstance(sha256, bytes):
                sha256 = sha256.decode('utf-8')

            rows.append({
                'sha256': sha256,
                's3_bucket': s3_bucket,
                's3_key': generate_s3_key_from_hash(sha256),
                'filetype_magika': row[1],
                'first_seen': row[2]  # From catalog_samples
            })

        date_info = f"from {start_date} to {end_date}"
        repo_info = f" for repository {repository}" if repository else ""
        print(f"[DEBUG] Found {len(rows)} objects in catalog_samples {date_info}{repo_info}")

        # Extract SHA256 hashes from the results for bulk checking
        sha256_list = [row.get('sha256') for row in rows if row.get('sha256')]

        if sha256_list and not force:
            # Check which hashes are already in the database
            existing_hashes = is_in_db_bulk(sha256_list, index_prefix, decompile, yara_scan, magika_filter)
            print(f"[INFO] Found {len(existing_hashes)} objects already in database")

            # Filter out rows with existing hashes
            rows = [row for row in rows if row.get('sha256') not in existing_hashes]
            print(f"[INFO] After filtering, {len(rows)} objects remain to be processed")
        elif force:
            print(f"[INFO] Force mode: skipping deduplication check, processing all {len(sha256_list)} objects")

        # Randomize the order of rows before returning
        import random
        random.shuffle(rows)

        return rows
    except Exception as e:
        print(f"[ERROR] Failed to query catalog_samples by date range: {e}")
        return []
    finally:
        if 'client' in locals():
            client.close()


def fetch_analyzed_samples(
    index_prefix: str,
    decompile: bool,
    magika_filter: Optional[str] = None,
    yara_scan: bool = False,
    force: bool = False,
    rerun: bool = False
) -> List[Dict]:
    """
    Query samples from basic_properties that have already been analyzed.
    Useful for reprocessing with decompilation or specific modules.

    Args:
        index_prefix: Table prefix for ClickHouse
        decompile: Whether we're in decompile mode (affects filtering)
        magika_filter: Optional single filetype to filter by (overrides SUPPORTED_FORMATS)
        yara_scan: Whether we're in YARA-only mode (skips decompile filtering)
        force: Skip all deduplication checks when True
        rerun: Query disassembled table directly (only already-disassembled samples)

    Returns:
        List of dicts with sha256, s3_bucket, s3_key for each analyzed sample
    """
    try:
        client = settings.create_clickhouse_client()
        s3_bucket = os.getenv('S3_BUCKET')

        if not s3_bucket:
            print("[ERROR] S3_BUCKET environment variable is required")
            return []

        # Rerun mode: query directly from disassembled_functions_references
        # instead of basic_properties. This targets only samples that already
        # went through binja successfully.
        if rerun:
            disassembled_table = _get_code_dedup_table(magika_filter)
            basic_table = f"{index_prefix}_basic_properties"

            # Join with basic_properties to get filetype_magika and apply format filters
            conditions = []
            params = {}

            if magika_filter:
                conditions.append("bp.filetype_magika = %(magika)s")
                params["magika"] = magika_filter
            else:
                supported_formats = get_supported_formats()
                conditions.append("bp.filetype_magika IN %(formats)s")
                params["formats"] = supported_formats

            query = (
                f"SELECT DISTINCT d.sha256, bp.filetype_magika "
                f"FROM {disassembled_table} d FINAL "
                f"INNER JOIN {basic_table} bp FINAL ON d.sha256 = bp.sha256"
            )
            if conditions:
                query += " WHERE " + " AND ".join(conditions)

            result = client.query(query, parameters=params)

            rows = []
            for row in result.result_rows:
                sha256 = row[0]
                if isinstance(sha256, bytes):
                    sha256 = sha256.decode('utf-8')
                rows.append({
                    'sha256': sha256,
                    's3_bucket': s3_bucket,
                    's3_key': generate_s3_key_from_hash(sha256),
                    'filetype_magika': row[1]
                })

            print(f"[INFO] Rerun mode: found {len(rows)} already-disassembled samples in {disassembled_table}")

            import random
            random.shuffle(rows)
            return rows

        basic_table = f"{index_prefix}_basic_properties"

        # Build query
        conditions = []
        params = {}

        if magika_filter:
            conditions.append("filetype_magika = %(magika)s")
            params["magika"] = magika_filter
        else:
            supported_formats = get_supported_formats()
            conditions.append("filetype_magika IN %(formats)s")
            params["formats"] = supported_formats

        query = f"SELECT DISTINCT sha256, filetype_magika FROM {basic_table} FINAL"
        if conditions:
            query += " WHERE " + " AND ".join(conditions)

        result = client.query(query, parameters=params)

        rows = []
        for row in result.result_rows:
            sha256 = row[0]
            if isinstance(sha256, bytes):
                sha256 = sha256.decode('utf-8')
            rows.append({
                'sha256': sha256,
                's3_bucket': s3_bucket,
                's3_key': generate_s3_key_from_hash(sha256),
                'filetype_magika': row[1]
            })

        print(f"[DEBUG] Found {len(rows)} analyzed samples in {basic_table}")

        # In YARA-only mode, filter out already-scanned samples (unless force)
        if yara_scan and not force:
            sha256_list = [r['sha256'] for r in rows]
            if sha256_list:
                already_scanned = _check_yara_matches_bulk(sha256_list)
                print(f"[INFO] Found {len(already_scanned)} already YARA-scanned samples")
                rows = [r for r in rows if r['sha256'].lower() not in already_scanned]
                print(f"[INFO] After filtering, {len(rows)} samples remain for YARA scanning")
        # In decompile mode, filter out already-disassembled samples (unless force)
        # We check disassembly (not decompilation) because disassembly is the ground truth:
        # disassembly always succeeds, decompilation may not, so a missing decompile
        # entry doesn't mean the sample wasn't analyzed.
        elif decompile and not force:
            sha256_list = [r['sha256'] for r in rows]
            if sha256_list:
                disassembled_table = _get_code_dedup_table(magika_filter)
                batch_size = 3900

                already_disassembled = set()
                for i in range(0, len(sha256_list), batch_size):
                    batch = sha256_list[i:i+batch_size]
                    placeholders = "','".join(batch)
                    check_query = f"SELECT DISTINCT sha256 FROM {disassembled_table} FINAL WHERE sha256 IN ('{placeholders}')"
                    check_result = client.query(check_query)
                    for check_row in check_result.result_rows:
                        hash_value = check_row[0]
                        if isinstance(hash_value, bytes):
                            hash_value = hash_value.decode('utf-8')
                        already_disassembled.add(hash_value)

                print(f"[INFO] Found {len(already_disassembled)} already-disassembled samples")
                rows = [r for r in rows if r['sha256'] not in already_disassembled]
                print(f"[INFO] After filtering, {len(rows)} samples remain for decompilation")
        elif force:
            print(f"[INFO] Force mode: returning all {len(rows)} analyzed samples")

        # Randomize the order of rows before returning
        import random
        random.shuffle(rows)

        return rows
    except Exception as e:
        print(f"[ERROR] Failed to query analyzed samples: {e}")
        return []
    finally:
        if 'client' in locals():
            client.close()


def is_in_db(sha256, index_prefix, client=None):
    """
    Check if a sample with given SHA256 already exists in the database

    Args:
        sha256: File's SHA256 hash
        index_prefix: Table prefix for ClickHouse
        client: Optional ClickHouse client instance

    Returns:
        bool: True if file exists, False otherwise
    """
    table = f"{index_prefix}_basic_properties"
    close_client = False

    try:
        if client is None:
            client = settings.create_clickhouse_client()
            close_client = True

        # No need for FINAL when checking existence with LIMIT 1
        # Any version of the row proves the sample exists
        query = f"SELECT 1 FROM {table} WHERE sha256 = %(sha256)s LIMIT 1"
        result = client.query(query, parameters={"sha256": sha256})

        return len(result.result_rows) > 0
    except Exception as e:
        print(f"[ERROR] Failed to check if file is in DB: {e}")
        return False
    finally:
        if close_client and client:
            client.close()


def _get_code_dedup_table(magika_filter: Optional[str] = None) -> str:
    """
    Return the code-analysis table used for deduplication based on filetype.

    APK samples are disassembled into code_apk_smali_methods_references;
    everything else uses {CLICKHOUSE_CODE_PREFIX}_disassembled_functions_references.
    """
    if magika_filter == "apk":
        return "code_apk_smali_methods_references"
    code_prefix = os.getenv("CLICKHOUSE_CODE_PREFIX", "code_binja")
    return f"{code_prefix}_disassembled_functions_references"


def is_in_code_db(sha256, client=None, filetype=None):
    """
    Check if a sample with given SHA256 has already been disassembled.

    Args:
        sha256: File's SHA256 hash
        client: Optional ClickHouse client instance
        filetype: Magika filetype label (e.g. 'apk') to pick the right code table

    Returns:
        bool: True if file has been disassembled, False otherwise
    """
    table = _get_code_dedup_table(filetype)
    close_client = False

    try:
        if client is None:
            client = settings.create_clickhouse_client()
            close_client = True

        query = f"SELECT 1 FROM {table} WHERE sha256 = %(sha256)s LIMIT 1"
        result = client.query(query, parameters={"sha256": sha256})

        return len(result.result_rows) > 0
    except Exception as e:
        print(f"[ERROR] Failed to check if file is in code DB: {e}")
        return False
    finally:
        if close_client and client:
            client.close()


def _check_yara_matches_bulk(sha256_list):
    """
    Check which samples have already been YARA-scanned by querying yara_matches.

    The yara_matches table uses FixedString(32) binary sha256, so we convert
    hex strings with unhex().

    Args:
        sha256_list: List of hex SHA256 strings to check

    Returns:
        set: Set of hex SHA256 hashes that already have YARA matches
    """
    # unhex('64hexchars') adds 9 chars overhead per entry vs plain '64hexchars'.
    # 3900 works for plain strings (~67 chars each = 261K) but overflows
    # max_query_size (262144) with unhex() wrapping (~73 chars each = 285K).
    # 3500 × 73 = 255K stays safely under the limit.
    batch_size = 3500
    already_scanned = set()

    try:
        client = settings.create_clickhouse_client()

        total_batches = (len(sha256_list) + batch_size - 1) // batch_size
        for batch_num, i in enumerate(range(0, len(sha256_list), batch_size), 1):
            batch = sha256_list[i:i+batch_size]
            unhex_list = ",".join(f"unhex('{h}')" for h in batch)
            query = f"SELECT DISTINCT hex(sha256) FROM yara_matches FINAL WHERE sha256 IN ({unhex_list})"
            result = client.query(query)
            for row in result.result_rows:
                hash_value = row[0]
                if isinstance(hash_value, bytes):
                    hash_value = hash_value.decode('utf-8')
                already_scanned.add(hash_value.lower())

            if batch_num % 100 == 0 or batch_num == total_batches:
                print(f"[INFO] YARA dedup progress: batch {batch_num}/{total_batches}, found {len(already_scanned)} so far")

        print(f"[INFO] YARA dedup: found {len(already_scanned)} already-scanned samples")
        return already_scanned

    except Exception as e:
        raise RuntimeError(f"YARA dedup query failed — aborting to prevent reprocessing all samples: {e}")
    finally:
        if 'client' in locals():
            client.close()


def is_in_db_bulk(sha256_list, index_prefix, decompile, yara_scan=False, magika_filter=None):
    """
    Check which samples from a list of SHA256 hashes should be skipped.

    Args:
        sha256_list: List of SHA256 hashes to check
        index_prefix: Table prefix for ClickHouse
        decompile: Whether we're in decompile mode
        yara_scan: Whether we're in YARA-only mode (checks yara_matches table)
        magika_filter: Magika filetype label (e.g. 'apk') to pick the right code table

    Returns:
        set: Set of SHA256 hashes that should be skipped
    """
    if not sha256_list:
        return set()

    # YARA-only mode: check yara_matches table for already-scanned samples
    if yara_scan:
        return _check_yara_matches_bulk(sha256_list)

    basic_table = f"{index_prefix}_basic_properties"
    batch_size = 3900

    try:
        client = settings.create_clickhouse_client()

        # Step 1: Check basic_properties (required for both modes)
        in_basic_properties = set()
        for i in range(0, len(sha256_list), batch_size):
            batch = sha256_list[i:i+batch_size]
            placeholders = "','".join(batch)
            query = f"SELECT DISTINCT sha256 FROM {basic_table} FINAL WHERE sha256 IN ('{placeholders}')"
            result = client.query(query)
            for row in result.result_rows:
                hash_value = row[0]
                if isinstance(hash_value, bytes):
                    hash_value = hash_value.decode('utf-8')
                in_basic_properties.add(hash_value)

        if not decompile:
            # Analysis mode: skip samples already in basic_properties
            return in_basic_properties

        # Decompile mode: skip samples NOT in basic_properties + already decompiled
        not_in_basic = set(sha256_list) - in_basic_properties

        if not in_basic_properties:
            return set(sha256_list)  # None ready for decompilation

        # Step 2: Check disassembled table for samples that ARE in basic_properties
        # Disassembly is the ground truth for code analysis (it always succeeds,
        # unlike decompilation which may fail).
        disassembled_table = _get_code_dedup_table(magika_filter)

        already_disassembled = set()
        samples_to_check = list(in_basic_properties)
        for i in range(0, len(samples_to_check), batch_size):
            batch = samples_to_check[i:i+batch_size]
            placeholders = "','".join(batch)
            query = f"SELECT DISTINCT sha256 FROM {disassembled_table} FINAL WHERE sha256 IN ('{placeholders}')"
            result = client.query(query)
            for row in result.result_rows:
                hash_value = row[0]
                if isinstance(hash_value, bytes):
                    hash_value = hash_value.decode('utf-8')
                already_disassembled.add(hash_value)

        return not_in_basic | already_disassembled

    except Exception as e:
        print(f"[ERROR] Failed to check hashes in DB: {e}")
        return set()
    finally:
        if 'client' in locals():
            client.close()