#!/usr/bin/env python

from __future__ import print_function

import os
import sys
import subprocess
import tempfile
import stat
import time
from datetime import datetime

LOG_DIR = "/var/duckdb"

if not os.path.exists(LOG_DIR):
    os.makedirs(LOG_DIR)

LOCK_FILE = "/var/duckdb/logs-ingest.lock"

def acquire_lock():
    if os.path.exists(LOCK_FILE):
        try:
            with open(LOCK_FILE, "r") as f:
                pid = int(f.read().strip())
            os.kill(pid, 0)
            print("[-] Error: Another instance is already running (PID %d). Exiting." % pid)
            sys.exit(0)
        except (ValueError, OSError):
            os.remove(LOCK_FILE)
    with open(LOCK_FILE, "w") as f:
        f.write(str(os.getpid()))

def release_lock():
    if os.path.exists(LOCK_FILE):
        os.remove(LOCK_FILE)

acquire_lock()

print("[+] Starting log ingestion process...")
start_time = time.time()

CURRENT_YEAR    = datetime.now().year
DB_FIREWALL     = os.path.join(LOG_DIR, "firewall.db")
DB_MESSAGES     = os.path.join(LOG_DIR, "messages.db")

OFFSET_FILE     = os.path.join(LOG_DIR, "firewall.offset")
TMP_LOG         = os.path.join(LOG_DIR, "firewall_delta.log")
MSG_OFFSET_FILE = os.path.join(LOG_DIR, "messages.offset")
TMP_MSG_LOG     = os.path.join(LOG_DIR, "messages_delta.log")

GEOIP_BASE = "/usr/lib/geoip"

def find_geoip_csv_dir(base):
    candidates = []
    for dirpath, dirnames, _ in os.walk(base):
        for d in dirnames:
            if d.startswith("GeoLite2-Country-CSV_"):
                candidates.append(os.path.join(dirpath, d))
    return sorted(candidates, reverse=True)[0] if candidates else None

GEOIP_CSV_DIR = find_geoip_csv_dir(GEOIP_BASE)
GEOIP_BLOCKS_V4 = None
GEOIP_LOCATIONS = None
GEOIP_VERSION_TAG = "NONE"

if not GEOIP_CSV_DIR:
    print("[!] Warning: GeoLite2-Country-CSV directory not found under %s. Geolocation will be skipped." % GEOIP_BASE)
else:
    GEOIP_BLOCKS_V4 = os.path.join(GEOIP_CSV_DIR, "GeoLite2-Country-Blocks-IPv4.csv")
    GEOIP_LOCATIONS = os.path.join(GEOIP_CSV_DIR, "GeoLite2-Country-Locations-en.csv")
    GEOIP_VERSION_TAG = os.path.basename(GEOIP_CSV_DIR)
    print("[+] MaxMind GeoIP database identified: %s" % GEOIP_VERSION_TAG)

if not os.path.exists(DB_FIREWALL) and os.path.exists(OFFSET_FILE):
    print("[!] Database firewall.db not found. Resetting firewall offset...")
    os.remove(OFFSET_FILE)

if not os.path.exists(DB_MESSAGES) and os.path.exists(MSG_OFFSET_FILE):
    print("[!] Database messages.db not found. Resetting messages offset...")
    os.remove(MSG_OFFSET_FILE)

DUCKDB_BIN = "/usr/bin/duckdb"
if not os.path.isfile(DUCKDB_BIN) or not os.access(DUCKDB_BIN, os.X_OK):
    print("[-] Critical Error: duckdb not found at /usr/bin/duckdb", file=sys.stderr)
    release_lock()
    sys.exit(1)

DUCKDB_THREADS = max(1, os.cpu_count() // 2) if hasattr(os, "cpu_count") else 1
if not hasattr(os, "cpu_count"):
    try:
        import multiprocessing
        DUCKDB_THREADS = max(1, multiprocessing.cpu_count() // 2)
    except Exception:
        DUCKDB_THREADS = 1

def run_logtail(log_file, offset_file, output_file):
    with open(output_file, "w") as out:
        with open(os.devnull, "w") as devnull:
            try:
                subprocess.call(
                    ["logtail", "-f", log_file, "-o", offset_file],
                    stdout=out,
                    stderr=devnull,
                )
            except OSError as e:
                print("[-] Critical Error: 'logtail' binary not found or failed to execute (%s). Please verify it is installed and available in PATH." % e, file=sys.stderr)
                release_lock()
                sys.exit(1)

print("[+] Collecting new log entries using logtail...")
run_logtail("/var/log/firewall", OFFSET_FILE,     TMP_LOG)
run_logtail("/var/log/messages", MSG_OFFSET_FILE, TMP_MSG_LOG)

fw_size = os.path.getsize(TMP_LOG) if os.path.exists(TMP_LOG) else 0
msg_size = os.path.getsize(TMP_MSG_LOG) if os.path.exists(TMP_MSG_LOG) else 0

if fw_size == 0 and msg_size == 0:
    print("[+] No new logs detected since last execution. Nothing to do.")
    for f in (TMP_LOG, TMP_MSG_LOG):
        if os.path.exists(f):
            os.remove(f)
    release_lock()
    sys.exit(0)

fw_status = 0
if fw_size > 0:
    print("[+] Importing Firewall log delta (%d bytes) into firewall.db..." % fw_size)
    SQL_FIREWALL = """
    PRAGMA threads=%(DUCKDB_THREADS)s;

    CREATE TABLE IF NOT EXISTS log_firewall (
        event_date TIMESTAMP,
        src_ip VARCHAR,
        dst_ip VARCHAR,
        dst_port INTEGER,
        interface VARCHAR,
        protocol VARCHAR,
        action VARCHAR,
        chain VARCHAR,
        src_country_code VARCHAR,
        src_country_name VARCHAR,
        dst_country_code VARCHAR,
        dst_country_name VARCHAR
    );

    ALTER TABLE log_firewall ADD COLUMN IF NOT EXISTS src_country_code VARCHAR;
    ALTER TABLE log_firewall ADD COLUMN IF NOT EXISTS src_country_name VARCHAR;
    ALTER TABLE log_firewall ADD COLUMN IF NOT EXISTS dst_country_code VARCHAR;
    ALTER TABLE log_firewall ADD COLUMN IF NOT EXISTS dst_country_name VARCHAR;

    WITH fw_lines AS (
        SELECT unnest(string_split_regex(content, '\\r?\\n')) AS line
        FROM read_text('%(TMP_LOG)s')
    ),
    new_fw AS (
        SELECT
            try_strptime('%(CURRENT_YEAR)s ' || substring(line, 1, 15), '%%Y %%b %%d %%H:%%M:%%S') AS event_date,
            string_split(string_split(line, 'SRC=')[2], ' ')[1] AS src_ip,
            string_split(string_split(line, 'DST=')[2], ' ')[1] AS dst_ip,
            CASE
                WHEN line LIKE '%%DPT=%%' THEN
                    TRY_CAST(string_split(string_split(line, 'DPT=')[2], ' ')[1] AS INTEGER)
                ELSE NULL
            END AS dst_port,
            CASE WHEN line LIKE '%%IN=%%'    THEN string_split(string_split(line, 'IN=')[2],    ' ')[1] ELSE NULL END AS interface,
            CASE WHEN line LIKE '%%PROTO=%%' THEN string_split(string_split(line, 'PROTO=')[2], ' ')[1] ELSE NULL END AS protocol,
            UPPER(COALESCE(
                regexp_extract(line, 'ulogd\\[\\d+\\]:\\s*([A-Z]+):([A-Z]+)', 2), 
                regexp_extract(line, 'ulogd\\[\\d+\\]:\\s*([A-Z]+):([A-Z]+)', 1), 
                'DROP'
            )) AS action,
            regexp_extract(line, 'ulogd\\[\\d+\\]:\\s*([A-Z]+)', 1) AS chain
        FROM fw_lines
        WHERE line IS NOT NULL
          AND line != ''
          AND line LIKE '%%SRC=%%'
    )
    INSERT INTO log_firewall (
        event_date, src_ip, dst_ip, dst_port, interface, protocol, action, chain,
        src_country_code, src_country_name, dst_country_code, dst_country_name
    )
    SELECT DISTINCT
        n.event_date, n.src_ip, n.dst_ip, n.dst_port, n.interface, n.protocol, n.action, n.chain,
        'PENDING' AS src_country_code, 'Processing...' AS src_country_name,
        'PENDING' AS dst_country_code, 'Processing...' AS dst_country_name
    FROM new_fw n
    WHERE n.event_date IS NOT NULL;
    """ % {
        "TMP_LOG":        TMP_LOG,
        "CURRENT_YEAR":   CURRENT_YEAR,
        "DUCKDB_THREADS": DUCKDB_THREADS,
    }

    tmp_sql_fd, tmp_sql_path = tempfile.mkstemp(suffix=".sql")
    try:
        with os.fdopen(tmp_sql_fd, "w") as tmp_sql_file:
            tmp_sql_file.write(SQL_FIREWALL)
        with open(tmp_sql_path, "r") as sql_input:
            proc = subprocess.Popen([DUCKDB_BIN, DB_FIREWALL], stdin=sql_input, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            _, stderr_out = proc.communicate()
            fw_status = proc.returncode
    finally:
        if os.path.exists(tmp_sql_path):
            os.remove(tmp_sql_path)
    if fw_status != 0:
        sys.stderr.write("[-] Firewall ingestion error:\n" + stderr_out.decode("utf-8", errors="replace"))
else:
    print("[+] No new Firewall logs to process.")

msg_status = 0
if msg_size > 0:
    print("[+] Importing OS log delta (%d bytes) into messages.db..." % msg_size)
    SQL_MESSAGES = """
    PRAGMA threads=%(DUCKDB_THREADS)s;

    CREATE TABLE IF NOT EXISTS log_messages (
        event_date TIMESTAMP,
        service VARCHAR,
        message VARCHAR
    );

    CREATE OR REPLACE TEMP TABLE temp_msg_lines AS
    SELECT unnest(string_split_regex(content, '\\r?\\n')) AS line
    FROM read_text('%(TMP_MSG_LOG)s')
    WHERE content IS NOT NULL AND content != '';

    INSERT INTO log_messages (event_date, service, message)
    SELECT
        strptime('%(CURRENT_YEAR)s ' || regexp_replace(substring(line, 1, 15), '\\s+', ' ', 'g'), '%%Y %%b %%d %%H:%%M:%%S') AS event_date,
        regexp_extract(line, '^[A-Za-z]{3}\\s+\\d+\\s+\\d{2}:\\d{2}:\\d{2}\\s+[^\\s]+\\s+([^:\\[]+)', 1) AS service,
        regexp_extract(line, '^[A-Za-z]{3}\\s+\\d+\\s+\\d{2}:\\d{2}:\\d{2}\\s+[^\\s]+\\s+[^:]+:\\s*(.*)$', 1) AS message
    FROM temp_msg_lines
    WHERE line IS NOT NULL
      AND line != ''
      AND length(line) > 16;
    """ % {
        "TMP_MSG_LOG":    TMP_MSG_LOG,
        "CURRENT_YEAR":   CURRENT_YEAR,
        "DUCKDB_THREADS": DUCKDB_THREADS,
    }

    tmp_sql_fd, tmp_sql_path = tempfile.mkstemp(suffix=".sql")
    try:
        with os.fdopen(tmp_sql_fd, "w") as tmp_sql_file:
            tmp_sql_file.write(SQL_MESSAGES)
        with open(tmp_sql_path, "r") as sql_input:
            proc = subprocess.Popen([DUCKDB_BIN, DB_MESSAGES], stdin=sql_input, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            _, stderr_out = proc.communicate()
            msg_status = proc.returncode
    finally:
        if os.path.exists(tmp_sql_path):
            os.remove(tmp_sql_path)
    if msg_status != 0:
        sys.stderr.write("[-] Messages ingestion error:\n" + stderr_out.decode("utf-8", errors="replace"))
else:
    print("[+] No new system logs (messages) to process.")

for f in (TMP_LOG, TMP_MSG_LOG):
    if os.path.exists(f):
        os.remove(f)

if fw_status == 0 and fw_size > 0 and GEOIP_CSV_DIR:
    print("[+] Checking GeoIP cache status in firewall.db...")
    must_reload = True
    try:
        proc_chk = subprocess.Popen([DUCKDB_BIN, DB_FIREWALL, "SELECT version FROM geoip_metadata;"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout_chk, _ = proc_chk.communicate()
        if GEOIP_VERSION_TAG in stdout_chk.decode('utf-8'):
            proc_cnt = subprocess.Popen([DUCKDB_BIN, DB_FIREWALL, "SELECT COUNT(*) FROM geoip_cache;"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stdout_cnt, _ = proc_cnt.communicate()
            if b"0" not in stdout_cnt:
                must_reload = False
    except Exception:
        pass

    SQL_GEOIP = "PRAGMA threads=%d;\n" % DUCKDB_THREADS
    SQL_GEOIP += "CREATE TABLE IF NOT EXISTS geoip_metadata (version VARCHAR);\n"
    SQL_GEOIP += "CREATE TABLE IF NOT EXISTS geoip_cache (ip_start BIGINT, ip_end BIGINT, country_code VARCHAR, country_name VARCHAR);\n"
    SQL_GEOIP += "UPDATE log_firewall SET src_country_code = 'PENDING', src_country_name = 'Processing...' WHERE src_country_code = 'LAN' AND src_ip NOT SIMILAR TO '(10\\..*|192\\.168\\..*|172\\.(1[6-9]|2[0-9]|3[0-1])\\..*|127\\..*)';\n"
    SQL_GEOIP += "UPDATE log_firewall SET dst_country_code = 'PENDING', dst_country_name = 'Processing...' WHERE dst_country_code = 'LAN' AND dst_ip NOT SIMILAR TO '(10\\..*|192\\.168\\..*|172\\.(1[6-9]|2[0-9]|3[0-1])\\..*|127\\..*)';\n"

    if must_reload:
        print("[!] Loading and indexing MaxMind GeoIP CSVs...")
        SQL_GEOIP += "DROP TABLE IF EXISTS geoip_cache;\n"
        SQL_GEOIP += """
        CREATE TABLE geoip_cache AS 
        SELECT
            ( Prague.ip_parts[1]::BIGINT << 24 ) + ( Prague.ip_parts[2]::BIGINT << 16 ) + ( Prague.ip_parts[3]::BIGINT << 8 ) + Prague.ip_parts[4]::BIGINT AS ip_start,
            (( Prague.ip_parts[1]::BIGINT << 24 ) + ( Prague.ip_parts[2]::BIGINT << 16 ) + ( Prague.ip_parts[3]::BIGINT << 8 ) + Prague.ip_parts[4]::BIGINT) + (power(2, 32 - Prague.mask) - 1)::BIGINT AS ip_end,
            loc.country_iso_code AS country_code,
            loc.country_name     AS country_name
        FROM (
            SELECT 
                string_to_array(string_split(network, '/')[1], '.') AS ip_parts,
                string_split(network, '/')[2]::INTEGER AS mask,
                registered_country_geoname_id,
                geoname_id
            FROM read_csv('%(GEOIP_BLOCKS_V4)s', AUTO_DETECT=TRUE)
        ) AS Prague
        JOIN read_csv('%(GEOIP_LOCATIONS)s', AUTO_DETECT=TRUE) AS loc
          ON COALESCE(NULLIF(Prague.geoname_id::VARCHAR, ''), Prague.registered_country_geoname_id::VARCHAR) = loc.geoname_id::VARCHAR
        WHERE loc.country_iso_code IS NOT NULL AND loc.country_iso_code != ''
        ORDER BY ip_start, ip_end;\n""" % {"GEOIP_BLOCKS_V4": GEOIP_BLOCKS_V4, "GEOIP_LOCATIONS": GEOIP_LOCATIONS}
        SQL_GEOIP += "DELETE FROM geoip_metadata; INSERT INTO geoip_metadata VALUES ('%s');\n" % GEOIP_VERSION_TAG
    else:
        print("[+] GeoIP cache is up to date.")

    print("[+] Running fast IP lookup match (BIGINT)...")
    SQL_GEOIP += """
    CREATE OR REPLACE TEMP TABLE pending_ips AS
    SELECT DISTINCT ip,
        (parts[1]::BIGINT << 24) + (parts[2]::BIGINT << 16) + (parts[3]::BIGINT << 8) + parts[4]::BIGINT AS ip_num
    FROM (
        SELECT ip, string_to_array(ip, '.') AS parts
        FROM (
            SELECT src_ip AS ip FROM log_firewall WHERE src_country_code = 'PENDING'
            UNION
            SELECT dst_ip AS ip FROM log_firewall WHERE dst_country_code = 'PENDING'
        )
        WHERE ip SIMILAR TO '[0-9]+\\.[0-9]+\\.[0-9]+\\.[0-9]+'
    ) AS parsed;

    CREATE OR REPLACE TEMP TABLE resolved_ips AS
    SELECT ip, country_code, country_name FROM (
        SELECT p.ip, geo.country_code, geo.country_name,
               ROW_NUMBER() OVER (PARTITION BY p.ip ORDER BY (geo.ip_end - geo.ip_start) ASC) AS rn
        FROM pending_ips AS p
        JOIN geoip_cache AS geo ON p.ip_num >= geo.ip_start AND p.ip_num <= geo.ip_end
    ) WHERE rn = 1;

    UPDATE log_firewall AS fw SET src_country_code = r.country_code, src_country_name = r.country_name FROM resolved_ips AS r WHERE fw.src_country_code = 'PENDING' AND fw.src_ip = r.ip;
    UPDATE log_firewall AS fw SET dst_country_code = r.country_code, dst_country_name = r.country_name FROM resolved_ips AS r WHERE fw.dst_country_code = 'PENDING' AND fw.dst_ip = r.ip;
    UPDATE log_firewall SET src_country_code = 'LAN', src_country_name = 'Local Network' WHERE src_country_code = 'PENDING';
    UPDATE log_firewall SET dst_country_code = 'LAN', dst_country_name = 'Local Network' WHERE dst_country_code = 'PENDING';
    """

    tmp_geo_fd, tmp_geo_path = tempfile.mkstemp(suffix=".sql")
    try:
        with os.fdopen(tmp_geo_fd, "w") as tmp_geo_file:
            tmp_geo_file.write(SQL_GEOIP)
        with open(tmp_geo_path, "r") as geo_input:
            proc_geo = subprocess.Popen([DUCKDB_BIN, DB_FIREWALL], stdin=geo_input, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            _, stderr_geo = proc_geo.communicate()
            geo_status = proc_geo.returncode
    finally:
        if os.path.exists(tmp_geo_path):
            os.remove(tmp_geo_path)

    if geo_status != 0:
        sys.stderr.write("[-] GeoIP enrichment error:\n" + stderr_geo.decode("utf-8", errors="replace"))
    else:
        print("[+] IP Geolocation enrichment completed successfully.")

print("[+] Executing data retention cleanup policy (30 days threshold)...")
if fw_status == 0 and os.path.exists(DB_FIREWALL):
    p_clean = subprocess.Popen([DUCKDB_BIN, DB_FIREWALL], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p_clean.communicate(input=b"PRAGMA threads=%d; DELETE FROM log_firewall WHERE event_date < CURRENT_DATE - INTERVAL 30 DAY;" % DUCKDB_THREADS)

if msg_status == 0 and os.path.exists(DB_MESSAGES):
    p_clean = subprocess.Popen([DUCKDB_BIN, DB_MESSAGES], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p_clean.communicate(input=b"PRAGMA threads=%d; DELETE FROM log_messages WHERE event_date < CURRENT_DATE - INTERVAL 30 DAY;" % DUCKDB_THREADS)

import grp
import pwd
for db_f in (DB_FIREWALL, DB_MESSAGES):
    if os.path.exists(db_f):
        try:
            uid = pwd.getpwnam("root").pw_uid
            gid = grp.getgrnam("nogroup").gr_gid
            os.chown(db_f, uid, gid)
        except KeyError:
            pass
        os.chmod(db_f, stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IWGRP | stat.S_IROTH)

release_lock()

end_time = time.time()
print("[+] ETL processing completed successfully in %.2f seconds!" % (end_time - start_time))
print("[+] Databases updated independently: firewall.db and messages.db.")