#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import csv
import subprocess
import datetime
import hashlib

os.environ["LD_LIBRARY_PATH"] = os.environ.get("LD_LIBRARY_PATH", "") + ":/usr/lib:/usr/local/lib:/lib"
os.environ["PATH"] = os.environ.get("PATH", "") + ":/usr/sbin:/sbin:/usr/local/sbin:/usr/bin:/bin"

GEOIP_DIR    = "/usr/lib/geoip"
EXTRACT_DIR  = os.path.join(GEOIP_DIR, "maxmind_deco")
RESTORE_DIR  = os.path.join(GEOIP_DIR, "ipset_global_restore")
PERSIST_FILE = os.path.join(GEOIP_DIR, "ipset-geo-persistent.save")
META_FILE          = os.path.join(GEOIP_DIR, "ipset-geo-persistent.meta")
INGEST_STATUS_FILE = os.path.join(GEOIP_DIR, "last_ingest")

MIN_GLOBAL_SUBNETS = 350000

# Complete ISO 3166-1 alpha-2 country list used as hardcoded fallback
ALL_COUNTRY_CODES = [
    "ad","ae","af","ag","ai","al","am","ao","aq","ar","as","at","au","aw","ax","az",
    "ba","bb","bd","be","bf","bg","bh","bi","bj","bl","bm","bn","bo","bq","br","bs",
    "bt","bv","bw","by","bz","ca","cc","cd","cf","cg","ch","ci","ck","cl","cm","cn",
    "co","cr","cu","cv","cw","cx","cy","cz","de","dj","dk","dm","do","dz","ec","ee",
    "eg","eh","er","es","et","fi","fj","fk","fm","fo","fr","ga","gb","gd","ge","gf",
    "gg","gh","gi","gl","gm","gn","gp","gq","gr","gs","gt","gu","gw","gy","hk","hm",
    "hn","hr","ht","hu","id","ie","il","im","in","io","iq","ir","is","it","je","jm",
    "jo","jp","ke","kg","kh","ki","km","kn","kp","kr","kw","ky","kz","la","lb","lc",
    "li","lk","lr","ls","lt","lu","lv","ly","ma","mc","md","me","mf","mg","mh","mk",
    "ml","mm","mn","mo","mp","mq","mr","ms","mt","mu","mv","mw","mx","my","mz","na",
    "nc","ne","nf","ng","ni","nl","no","np","nr","nu","nz","om","pa","pe","pf","pg",
    "ph","pk","pl","pm","pn","pr","ps","pt","pw","py","qa","re","ro","rs","ru","rw",
    "sa","sb","sc","sd","se","sg","sh","si","sj","sk","sl","sm","sn","so","sr","ss",
    "st","sv","sx","sy","sz","tc","td","tf","tg","th","tj","tk","tl","tm","tn","to",
    "tr","tt","tv","tw","tz","ua","ug","um","us","uy","uz","va","vc","ve","vg","vi",
    "vn","vu","wf","ws","ye","yt","za","zm","zw"
]


def log_message(message):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    formatted_msg = "[%s] %s" % (timestamp, message)
    print(formatted_msg)
    try:
        cmd = "logger -t 'geoip-ingestor[%d]' -p daemon.info -- '%s'" % (os.getpid(), message.replace("'", "'\\''"))
        subprocess.call(cmd, shell=True)
    except Exception as e:
        sys.stderr.write("[-] CRITICAL LOG ERROR: Cannot write to syslog-ng via logger: %s\n" % e)

def init_log_file():
    try:
        cmd = "logger -t 'geoip-ingestor[%d]' -p daemon.info -- '=== OpenFW GeoIP Ingestor — Session Init ==='" % os.getpid()
        subprocess.call(cmd, shell=True)
    except Exception as e:
        sys.stderr.write("[-] CRITICAL LOG ERROR: Cannot initialize syslog-ng session: %s\n" % e)



# ---------------------------------------------------------------------------
# Metadata helpers
# ---------------------------------------------------------------------------

def write_meta(country_counts):
    try:
        with open(META_FILE, 'w') as f:
            for iso, count in sorted(country_counts.items()):
                f.write("%s=%d\n" % (iso, count))
        log_message("[+] Metadata file saved: %s (%d entries)" % (META_FILE, len(country_counts)))
        return True
    except Exception as e:
        log_message("[-] WARNING: Could not write metadata file: %s" % e)
        return False

def load_meta():
    meta = {}
    if not os.path.exists(META_FILE):
        return meta
    try:
        with open(META_FILE, 'r') as f:
            for line in f:
                line = line.strip()
                if '=' in line and not line.startswith('#'):
                    iso, count = line.split('=', 1)
                    meta[iso.strip()] = int(count.strip())
    except Exception as e:
        log_message("[-] WARNING: Could not read metadata file: %s" % e)
    return meta


# ---------------------------------------------------------------------------
# Kernel health checks
# ---------------------------------------------------------------------------

def check_sets_healthy(prefix, meta):
    if not meta:
        return False, []

    unhealthy = []
    for iso, expected in meta.items():
        set_name = "%s_%s" % (prefix, iso)
        ret = subprocess.call("ipset list %s >/dev/null 2>&1" % set_name, shell=True)
        if ret != 0:
            log_message("[-] HEALTH: Set %s missing from kernel." % set_name)
            unhealthy.append(iso)
            continue
        try:
            out = subprocess.check_output("ipset list %s | wc -l" % set_name, shell=True)
            if int(out.strip()) < 10:
                log_message("[-] HEALTH: Set %s appears empty (wc -l < 10)." % set_name)
                unhealthy.append(iso)
        except Exception as e:
            log_message("[-] HEALTH: Could not query %s: %s" % (set_name, e))
            unhealthy.append(iso)

    all_healthy = (len(unhealthy) == 0)
    return all_healthy, unhealthy

def verify_boot_sets(prefix, persist_file):
    expected = set()
    try:
        with open(persist_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith("create "):
                    parts = line.split()
                    if len(parts) >= 2:
                        expected.add(parts[1])
    except Exception as e:
        log_message("[-] WARNING: Could not parse snapshot for verification: %s" % e)
        return False, [], []

    missing = []
    empty   = []
    for set_name in expected:
        ret = subprocess.call("ipset list %s >/dev/null 2>&1" % set_name, shell=True)
        if ret != 0:
            missing.append(set_name)
            continue
        try:
            out = subprocess.check_output("ipset list %s | wc -l" % set_name, shell=True)
            if int(out.strip()) < 10:
                empty.append(set_name)
        except Exception:
            empty.append(set_name)

    ok = (len(missing) == 0 and len(empty) == 0)
    return ok, missing, empty


# ---------------------------------------------------------------------------
# Fallback skeleton creators
# ---------------------------------------------------------------------------

def create_fallback_skeletons_from_all_countries(prefix):
    log_message("[*] Creating full skeleton for %d known ISO countries..." % len(ALL_COUNTRY_CODES))
    created = 0
    for iso in ALL_COUNTRY_CODES:
        set_name = "%s_%s" % (prefix, iso)
        maxelem  = 200000 if iso == "us" else 65536
        ret = subprocess.call("ipset list %s >/dev/null 2>&1" % set_name, shell=True)
        if ret != 0:
            subprocess.call("ipset create %s hash:net maxelem %d -!" % (set_name, maxelem), shell=True)
            created += 1
    log_message("[+] Full skeleton applied: %d sets created." % created)

def create_fallback_skeletons_from_snapshot(prefix, persist_file):
    countries_in_snapshot = {}
    try:
        with open(persist_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith("create "):
                    parts = line.split()
                    if len(parts) >= 5:
                        set_name = parts[1]
                        try:
                            maxelem_idx = parts.index("maxelem")
                            maxelem_val = int(parts[maxelem_idx + 1].rstrip(",").split()[0])
                        except (ValueError, IndexError):
                            maxelem_val = 65536
                        countries_in_snapshot[set_name] = maxelem_val
    except Exception as e:
        log_message("[-] WARNING: Could not parse snapshot for skeleton creation: %s" % e)

    if not countries_in_snapshot:
        log_message("[*] Snapshot unreadable -- falling back to full ISO country skeleton...")
        create_fallback_skeletons_from_all_countries(prefix)
        return

    log_message("[*] Creating fallback skeleton for %d sets found in snapshot..." % len(countries_in_snapshot))
    created = 0
    for set_name, maxelem in countries_in_snapshot.items():
        ret = subprocess.call("ipset list %s >/dev/null 2>&1" % set_name, shell=True)
        if ret != 0:
            subprocess.call("ipset create %s hash:net maxelem %d -!" % (set_name, maxelem), shell=True)
            created += 1
    log_message("[+] Skeleton applied: %d missing sets created (empty but present)." % created)


# ---------------------------------------------------------------------------
# Boot recovery
# ---------------------------------------------------------------------------

def execute_cold_start_boot(prefix):
    log_message("=== OpenFW: Starting Unified Cold Start Boot Process ===")

    meta = load_meta()
    if meta:
        log_message("[*] Metadata found (%d sets). Validating kernel state..." % len(meta))
        all_healthy, unhealthy = check_sets_healthy(prefix, meta)
        if all_healthy:
            log_message("[+] All sets are present and healthy in kernel. Skipping restore.")
            return True
        else:
            log_message("[*] %d set(s) unhealthy: %s. Proceeding with restore..." % (
                len(unhealthy), ", ".join(unhealthy)))
    else:
        log_message("[*] No metadata file found. Proceeding with full restore...")

    if os.path.exists(PERSIST_FILE) and os.path.getsize(PERSIST_FILE) > 0:
        log_message("[+] Found persistent GeoIP snapshot. Restoring all sets into Kernel...")
        try:
            subprocess.check_call("ipset restore < %s" % PERSIST_FILE, shell=True)
            log_message("[+] Kernel memory populated. Verifying integrity of restored sets...")

            ok, missing, empty = verify_boot_sets(prefix, PERSIST_FILE)
            if ok:
                log_message("[+] Integrity check passed. All sets present and populated.")
                return True
            else:
                if missing:
                    log_message("[-] WARNING: %d sets missing after restore: %s" % (len(missing), ", ".join(missing)))
                if empty:
                    log_message("[-] WARNING: %d sets restored empty: %s" % (len(empty), ", ".join(empty)))
                log_message("[*] Partial restore detected. Applying skeleton for missing sets...")
                create_fallback_skeletons_from_snapshot(prefix, PERSIST_FILE)
                return False

        except subprocess.CalledProcessError as e:
            log_message("[-] CRITICAL: ipset restore failed: %s" % e)
            log_message("[*] Snapshot may be corrupted. Applying full skeleton fallback...")
            create_fallback_skeletons_from_snapshot(prefix, PERSIST_FILE)
            return False

    else:
        log_message("[-] WARNING: No persistent snapshot found at %s" % PERSIST_FILE)

    log_message("[*] No snapshot available. Creating full skeleton from known ISO country list...")
    create_fallback_skeletons_from_all_countries(prefix)
    return False


# ---------------------------------------------------------------------------
# CSV parsing + staging file generation
# ---------------------------------------------------------------------------

def locate_csv_files():
    csv_blocks, csv_locations = None, None
    for root, _, files in os.walk(EXTRACT_DIR):
        for f in files:
            if "Blocks-IPv4" in f:
                csv_blocks = os.path.join(root, f)
            elif "Locations-en" in f or "Locations-pt-BR" in f:
                csv_locations = os.path.join(root, f)

    if not csv_blocks or not csv_locations:
        log_message("[-] CRITICAL: Missing vital CSV files inside extracted directory.")
        return None, None
    return csv_blocks, csv_locations


def load_country_map(csv_locations):
    country_map = {}
    try:
        mode   = 'r' if sys.version_info[0] >= 3 else 'rb'
        kwargs = {'encoding': 'utf-8'} if sys.version_info[0] >= 3 else {}
        with open(csv_locations, mode, **kwargs) as f:
            reader = csv.DictReader(f)
            for row in reader:
                geo_id   = row.get('geoname_id')
                iso_code = row.get('country_iso_code')
                if geo_id and iso_code:
                    country_map[geo_id] = iso_code.lower()
    except Exception as e:
        log_message("[-] Error building country map: %s" % e)
    return country_map

def generate_staging_files(csv_blocks, country_map, prefix):
    if os.path.exists(RESTORE_DIR):
        import shutil
        shutil.rmtree(RESTORE_DIR)
    os.makedirs(RESTORE_DIR)

    file_handlers          = {}
    country_counts         = {}
    total_processed_subnets = 0

    try:
        log_message("[+] Analyzing blocks and writing staging files to disk...")
        mode   = 'r' if sys.version_info[0] >= 3 else 'rb'
        kwargs = {'encoding': 'utf-8'} if sys.version_info[0] >= 3 else {}

        with open(csv_blocks, mode, **kwargs) as f_in:
            reader = csv.DictReader(f_in)
            if 'network' not in reader.fieldnames:
                log_message("[-] CRITICAL: Blocks CSV header missing 'network' column.")
                return None

            for row in reader:
                geo_id = row.get('geoname_id') or row.get('registered_country_geoname_id')
                if geo_id in country_map:
                    country_iso = country_map[geo_id]
                    cidr        = row.get('network')

                    if cidr and '.' in cidr:
                        if country_iso not in file_handlers:
                            restore_path = os.path.join(RESTORE_DIR, "%s.txt" % country_iso)
                            file_handlers[country_iso] = open(restore_path, 'w')
                            ipset_tmp = "%s_%s_tmp" % (prefix, country_iso)
                            max_elem  = 200000 if country_iso == "us" else 65536
                            file_handlers[country_iso].write("create %s hash:net maxelem %d -!\n" % (ipset_tmp, max_elem))
                            file_handlers[country_iso].write("flush %s\n" % ipset_tmp)
                            country_counts[country_iso] = 0

                        ipset_tmp = "%s_%s_tmp" % (prefix, country_iso)
                        file_handlers[country_iso].write("add %s %s -!\n" % (ipset_tmp, cidr))
                        country_counts[country_iso]   += 1
                        total_processed_subnets       += 1

        for fh in file_handlers.values():
            fh.close()

        log_message("[+] Global Parsing Summary: Total of %d valid IPv4 subnets found." % total_processed_subnets)

        if total_processed_subnets < MIN_GLOBAL_SUBNETS:
            log_message("[-] CRITICAL BREAK: Total subnets (%d) is below safety threshold (%d)." % (total_processed_subnets, MIN_GLOBAL_SUBNETS))
            return None

        return country_counts

    except Exception as e:
        log_message("[-] Error during staging file generation: %s" % e)
        return None
    finally:
        for fh in file_handlers.values():
            if not fh.closed:
                fh.close()


# ---------------------------------------------------------------------------
# Kernel ingestion
# ---------------------------------------------------------------------------

def apply_safely_to_kernel(country_counts, prefix):
    log_message("[+] Executing kernel ingestion loop with atomic verification...")
    success_countries = 0
    failed_countries  = 0

    for country_iso, expected_count in country_counts.items():
        if expected_count == 0:
            continue

        ipset_main   = "%s_%s" % (prefix, country_iso)
        ipset_tmp    = "%s_%s_tmp" % (prefix, country_iso)
        restore_path = os.path.join(RESTORE_DIR, "%s.txt" % country_iso)
        max_elem     = 200000 if country_iso == "us" else 65536

        try:
            subprocess.call("ipset destroy %s >/dev/null 2>&1" % ipset_tmp, shell=True)

            check_existing = subprocess.call("ipset list %s >/dev/null 2>&1" % ipset_main, shell=True)
            if check_existing != 0:
                subprocess.check_call("ipset create %s hash:net maxelem %d" % (ipset_main, max_elem), shell=True)

            subprocess.check_call("ipset restore < %s" % restore_path, shell=True)

            output = subprocess.check_output("ipset list %s | wc -l" % ipset_tmp, shell=True)
            if int(output.strip()) < 10:
                log_message("[-] WARNING: Inconsistency on %s_tmp. Expected ~%d entries but set appears empty. SWAP DENIED." % (ipset_tmp, expected_count))
                subprocess.call("ipset destroy %s >/dev/null 2>&1" % ipset_tmp, shell=True)
                failed_countries += 1
                continue

            subprocess.check_call("ipset swap %s %s" % (ipset_tmp, ipset_main), shell=True)
            subprocess.call("ipset destroy %s >/dev/null 2>&1" % ipset_tmp, shell=True)
            success_countries += 1

        except Exception as e:
            log_message("[-] WARNING: Failed to update country [%s]: %s. Previous set preserved." % (country_iso, e))
            subprocess.call("ipset destroy %s >/dev/null 2>&1" % ipset_tmp, shell=True)
            failed_countries += 1
            continue

    log_message("[+] Synchronization Finished. Updated: %d countries. Failed/Preserved: %d countries." % (success_countries, failed_countries))
    return True

# ---------------------------------------------------------------------------
# CSV change detection via MD5 Hash
# ---------------------------------------------------------------------------

def calculate_file_hash(filepath):
    """Calculates the MD5 hash of a file in chunks to be memory efficient."""
    hasher = hashlib.md5()
    try:
        with open(filepath, 'rb') as f:
            for chunk in iter(lambda: f.read(65536), b''):
                hasher.update(chunk)
        return hasher.hexdigest()
    except Exception as e:
        log_message("[-] Error calculating hash for %s: %s" % (filepath, e))
        return None

def csv_changed_since_last_ingest():
    """
    Compares the MD5 hash of the MaxMind CSV blocks file against the hash
    of the last successful ingest recorded in INGEST_STATUS_FILE.
    Returns True if the content changed or if no previous record exists.
    """
    # Find the blocks CSV to check its hash
    csv_blocks = None
    for root, _, files in os.walk(EXTRACT_DIR):
        for f in files:
            if "Blocks-IPv4" in f:
                csv_blocks = os.path.join(root, f)
                break
        if csv_blocks:
            break

    if not csv_blocks:
        log_message("[-] WARNING: Could not locate Blocks CSV to check hash. Forcing ingest.")
        return True

    current_hash = calculate_file_hash(csv_blocks)
    if not current_hash:
        log_message("[-] WARNING: Could not calculate current CSV hash. Forcing ingest.")
        return True

    if not os.path.exists(INGEST_STATUS_FILE):
        log_message("[*] No previous ingest record found. Ingest required.")
        return True

    try:
        with open(INGEST_STATUS_FILE, 'r') as f:
            last_ingest_hash = f.read().strip()
    except Exception as e:
        log_message("[-] WARNING: Could not read ingest status file: %s. Forcing ingest." % e)
        return True

    if current_hash != last_ingest_hash:
        log_message("[*] The local CSV (New Hash: %s, Old Hash: %s). Ingest required." % (current_hash, last_ingest_hash))
        return True

    log_message("[+] The local CSV is the same as the last injection! No changes...")
    return False

def update_ingest_timestamp():
    """Records the current blocks CSV MD5 hash as the last successful ingest marker."""
    csv_blocks = None
    for root, _, files in os.walk(EXTRACT_DIR):
        for f in files:
            if "Blocks-IPv4" in f:
                csv_blocks = os.path.join(root, f)
                break
        if csv_blocks:
            break

    if not csv_blocks:
        return

    current_hash = calculate_file_hash(csv_blocks)
    if current_hash:
        try:
            with open(INGEST_STATUS_FILE, 'w') as f:
                f.write(current_hash)
            log_message("[+] Ingest hash signature updated: %s" % current_hash)
        except Exception as e:
            log_message("[-] WARNING: Could not write ingest status file: %s" % e)


def save_snapshot(country_counts):
    try:
        log_message("[+] Saving active production datasets to persistent snapshot...")
        with open(PERSIST_FILE, 'w') as f_persist:
            for country_iso in country_counts.keys():
                txt_path = os.path.join(RESTORE_DIR, "%s.txt" % country_iso)
                if os.path.exists(txt_path):
                    with open(txt_path, 'r') as f_src:
                        for line in f_src:
                            f_persist.write(line.replace("_tmp", ""))
        log_message("[+] Snapshot updated successfully at: %s" % PERSIST_FILE)
        write_meta(country_counts)
        update_ingest_timestamp()
    except Exception as e:
        log_message("[-] WARNING: Sets updated but persistent backup failed: %s" % e)

def clean_staging():
    if os.path.exists(RESTORE_DIR):
        os.chdir("/")
        import shutil
        shutil.rmtree(RESTORE_DIR)


if __name__ == "__main__":
    if os.getuid() != 0:
        sys.stderr.write("[-] Error: Root privileges required.\n")
        sys.exit(1)

    parser = argparse.ArgumentParser(description="OpenFW GeoIP Ingestor — Kernel Ingestion & Boot Recovery Engine")
    parser.add_argument("--prefix",      default="geo", help="Prefix for ipset names in kernel")
    parser.add_argument("--boot",        action="store_true", help="Execute cold start restore from snapshot")
    parser.add_argument("--force",       action="store_true", help="Force ingestion even if CSV hash matches")
    args = parser.parse_args()

    init_log_file()

    if args.boot:
        if execute_cold_start_boot(args.prefix.lower()):
            log_message("=== OpenFW GeoIP Ingestor: Boot recovery executed successfully ===")
            sys.exit(0)
        else:
            log_message("[-] OpenFW GeoIP Ingestor: Boot recovery finished with limitations.")
            sys.exit(1)

    log_message("=== OpenFW GeoIP Ingestor: Starting BULLETPROOF Global Mode ===")

    if not os.path.exists(EXTRACT_DIR):
        log_message("[-] CRITICAL: EXTRACT_DIR not found at %s. Run geoip-updater.py first." % EXTRACT_DIR)
        sys.exit(1)

    # ---------------------------------------------------------------------------
    # CORREÇÃO DA ORDEM DE VALIDAÇÃO: Primeiro garante a saúde da memória do Kernel
    # ---------------------------------------------------------------------------
    meta = load_meta()
    kernel_needs_update = False

    if meta:
        log_message("[*] Validating kernel memory sets state before file checks...")
        all_healthy, unhealthy = check_sets_healthy(args.prefix.lower(), meta)
        if not all_healthy:
            log_message("[!] Kernel compromised! %d set(s) unhealthy/missing. Forcing full ingestion..." % len(unhealthy))
            kernel_needs_update = True
    else:
        log_message("[*] No metadata found. Forcing full ingestion to populate kernel...")
        kernel_needs_update = True

    # Se a memória estiver saudável E a flag --force NÃO foi usada, validamos o arquivo CSV
    if not kernel_needs_update and not args.force:
        if not csv_changed_since_last_ingest():
            log_message("[+] Execution skipped because files are equal and kernel is healthy.")
            sys.exit(0)
    elif args.force:
        log_message("[*] Force option detected. Bypassing validations and running ingestion...")
    elif kernel_needs_update:
        log_message("[*] Ingestion triggered due to unhealthy kernel sets.")

    # ---------------------------------------------------------------------------

    csv_blocks, csv_locations = locate_csv_files()
    if not csv_blocks or not csv_locations:
        sys.exit(1)

    country_map = load_country_map(csv_locations)
    if not country_map:
        sys.exit(1)

    country_counts = generate_staging_files(csv_blocks, country_map, args.prefix.lower())
    if not country_counts:
        clean_staging()
        sys.exit(1)

    if apply_safely_to_kernel(country_counts, args.prefix.lower()):
        log_message("[+] Global update complete. All countries synchronized successfully.")
        save_snapshot(country_counts)
        clean_staging()
    else:
        log_message("[-] Global application encountered critical errors.")
        clean_staging()
        sys.exit(1)

    log_message("=== OpenFW GeoIP Ingestor: Global run completed successfully ===")