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

import os
import sys
import argparse
import zipfile
import subprocess
import datetime
import pwd
import grp
import time

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"

CONFIG_FILE = "/var/efw/geoip/maxmind.conf"
STATUS_FILE = "/var/efw/geoip/last_update"
ERROR_FILE  = "/var/efw/geoip/last_error"

GEOIP_DIR   = "/usr/lib/geoip"
ZIP_FILE    = os.path.join(GEOIP_DIR, "GeoLite2-Country-CSV.zip")
EXTRACT_DIR = os.path.join(GEOIP_DIR, "maxmind_deco")

MIN_ZIP_SIZE       = 2000000
LAST_MODIFIED_FILE = ZIP_FILE + ".last_modified"


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-updater[%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-updater[%d]' -p daemon.info -- '=== OpenFW GeoIP Updater — 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)

def load_config():
    config = {}
    try:
        with open(CONFIG_FILE, 'r') as f:
            for line in f:
                line = line.strip()
                if '=' in line and not line.startswith('#'):
                    k, v = line.split('=', 1)
                    config[k.strip()] = v.strip()
    except Exception as e:
        log_message("[-] CRITICAL: Cannot read config file %s: %s" % (CONFIG_FILE, e))
        sys.exit(1)
    if not config.get('LICENSE_KEY'):
        log_message("[-] CRITICAL: LICENSE_KEY not found in config file.")
        sys.exit(1)
    return config

def write_status(success):
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    try:
        if success:
            with open(STATUS_FILE, 'w') as f:
                f.write(now)
            if os.path.exists(ERROR_FILE):
                os.remove(ERROR_FILE)
            log_message("[+] Status file updated: %s" % now)
        else:
            with open(ERROR_FILE, 'w') as f:
                f.write(now)
    except Exception as e:
        log_message("[-] WARNING: Could not write status file: %s" % e)

def download_maxmind_base(force=False):
    if force and os.path.exists(ZIP_FILE):
        log_message("[!] Force download requested. Invalidating local ZIP cache...")
        try:
            os.remove(ZIP_FILE)
        except Exception as e:
            log_message("[-] Warning: Could not clear local zip: %s" % e)

    if os.path.exists(ZIP_FILE) and os.path.getsize(ZIP_FILE) >= MIN_ZIP_SIZE:
        log_message("[*] MaxMind base already exists locally in %s. Skipping download." % GEOIP_DIR)
        return True

    log_message("[+] Downloading official MaxMind base via curl...")

    curl_cmd = "curl -f -sSL --connect-timeout 15 --retry 3 --max-redirs 5 --location-trusted -o %s '%s'" % (ZIP_FILE, MAXMIND_URL)

    try:
        subprocess.check_call(curl_cmd, shell=True)
    except subprocess.CalledProcessError as e:
        log_message("[-] CRITICAL: Curl download failed: %s" % e)
        return False

    if not os.path.exists(ZIP_FILE):
        log_message("[-] CRITICAL: Downloaded file is MISSING from disk.")
        return False

    actual_size = os.path.getsize(ZIP_FILE)
    if actual_size < MIN_ZIP_SIZE:
        log_message("[-] CRITICAL: Downloaded file is too small (%d bytes). Expected at least %d bytes. Check your MaxMind License Key!" % (actual_size, MIN_ZIP_SIZE))
        return False

    log_message("[+] Download completed and file size verified (%d bytes)." % actual_size)
    return True

def extract_and_validate_zip():
    try:
        if os.path.exists(EXTRACT_DIR):
            import shutil
            shutil.rmtree(EXTRACT_DIR)

        log_message("[+] Verifying ZIP structural integrity...")
        with zipfile.ZipFile(ZIP_FILE, 'r') as zip_ref:
            if zip_ref.testzip() is not None:
                log_message("[-] CRITICAL: Zip file is corrupted or truncated.")
                return False
            log_message("[+] ZIP file is healthy. Extracting data...")
            zip_ref.extractall(EXTRACT_DIR)
        return True
    except Exception as e:
        log_message("[-] Error manipulating ZIP: %s" % e)
        return False

def get_remote_last_modified(url_base):
    curl_cmd = "curl -f -sI --connect-timeout 15 --retry 3 --max-redirs 5 --location-trusted '%s'" % url_base
    try:
        result = subprocess.check_output(curl_cmd, shell=True)
        for line in result.splitlines():
            line = line.decode('utf-8', errors='replace').strip()
            if line.lower().startswith("last-modified:"):
                return line.split(":", 1)[1].strip()
        log_message("[-] WARNING: Last-Modified header not found in HEAD response.")
        return None
    except subprocess.CalledProcessError as e:
        log_message("[-] WARNING: Could not fetch HEAD from MaxMind: %s" % e)
        return None

def get_local_last_modified():
    if not os.path.exists(LAST_MODIFIED_FILE):
        return None
    try:
        with open(LAST_MODIFIED_FILE, 'r') as f:
            return f.read().strip()
    except Exception as e:
        log_message("[-] WARNING: Could not read local Last-Modified cache: %s" % e)
        return None

def save_local_last_modified(value):
    try:
        with open(LAST_MODIFIED_FILE, 'w') as f:
            f.write(value)
    except Exception as e:
        log_message("[-] WARNING: Could not save local Last-Modified cache: %s" % e)

def is_up_to_date(url_base):
    log_message("[*] Checking remote Last-Modified header to detect if update is available...")
    remote_lm = get_remote_last_modified(url_base)
    if remote_lm is None:
        log_message("[!] Could not retrieve remote Last-Modified. Proceeding with download to be safe.")
        return False
    local_lm = get_local_last_modified()
    if local_lm is None:
        log_message("[*] No local Last-Modified cache found. Proceeding with download.")
        return False
    if remote_lm == local_lm:
        log_message("[*] Remote Last-Modified matches local cache (%s). Database is already up to date." % remote_lm)
        return True
    log_message("[+] Update detected (local=%s, remote=%s)." % (local_lm, remote_lm))
    return False

def clean_zip():
    if os.path.exists(ZIP_FILE):
        try:
            os.remove(ZIP_FILE)
        except Exception as e:
            log_message("[-] WARNING: Could not remove ZIP file: %s" % e)

def update_timestamp_file():
    target_file = "/var/efw/download/timestamps"
    target_dir = os.path.dirname(target_file)
    current_timestamp = str(int(time.time()))
    line_to_write = "GEOIP_MAXMIND_DB_URL=%s\n" % current_timestamp
    
    try:
        uid = pwd.getpwnam("nobody").pw_uid
        gid = grp.getgrnam("nogroup").gr_gid

        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
            os.chown(target_dir, uid, gid)
            log_message("[+] Created directory %s and set to nobody:nogroup" % target_dir)
            
        lines = []
        found = False
        
        if os.path.exists(target_file):
            with open(target_file, 'r') as f:
                for line in f:
                    if line.startswith("GEOIP_MAXMIND_DB_URL="):
                        lines.append(line_to_write)
                        found = True
                    else:
                        lines.append(line)
        
        if not found:
            lines.append(line_to_write)
            
        with open(target_file, 'w') as f:
            f.writelines(lines)
            
        log_message("[+] Timestamps file updated with GEOIP_MAXMIND_DB_URL=%s" % current_timestamp)
        
        os.chown(target_file, uid, gid)
        log_message("[+] Permissions for %s set to nobody:nogroup" % target_file)
        return True
    except Exception as e:
        log_message("[-] Error updating timestamps file: %s" % e)
        return False


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

    parser = argparse.ArgumentParser(description="OpenFW GeoIP Updater — Download & Extraction Engine")
    parser.add_argument("--force-download", action="store_true", help="Bypass local ZIP cache and force re-download")
    parser.add_argument("--clear-cache",    action="store_true", help="Remove ZIP file after extraction")
    args = parser.parse_args()

    init_log_file()

    CONFIG      = load_config()
    LICENSE_KEY = CONFIG['LICENSE_KEY']
    EDITION     = CONFIG.get('EDITION', 'GeoLite2-Country-CSV')
    MAXMIND_URL = "https://download.maxmind.com/app/geoip_download?edition_id=%s&license_key=%s&suffix=zip" % (EDITION, LICENSE_KEY)

    log_message("=== OpenFW GeoIP Updater: Starting ===")

    if not os.path.exists(GEOIP_DIR):
        try:
            os.makedirs(GEOIP_DIR)
        except Exception as dir_err:
            log_message("[-] CRITICAL: Failed to create base directory %s: %s" % (GEOIP_DIR, dir_err))
            write_status(success=False)
            sys.exit(1)

    if not args.force_download and is_up_to_date(MAXMIND_URL):
        log_message("=== OpenFW GeoIP Updater: No update needed. Exiting. ===")
        sys.exit(0)

    if not download_maxmind_base(force=args.force_download):
        write_status(success=False)
        sys.exit(1)

    # Salva o Last-Modified remoto após download bem-sucedido
    remote_lm_after = get_remote_last_modified(MAXMIND_URL)
    if remote_lm_after:
        save_local_last_modified(remote_lm_after)

    if not extract_and_validate_zip():
        write_status(success=False)
        sys.exit(1)

    if args.clear_cache:
        clean_zip()
        log_message("[+] ZIP file cleared (--clear-cache).")

    update_timestamp_file()

    log_message("[+] CSVs ready in %s" % EXTRACT_DIR)
    write_status(success=True)
    log_message("=== OpenFW GeoIP Updater: Completed successfully ===")
    sys.exit(0)