#!/usr/bin/env python3
import sqlite3, subprocess, re, os, requests, time

DB = "sshd_failures.db"
STATE = "cursor.txt"

def get_cursor():
    return open(STATE).read().strip() if os.path.exists(STATE) else None

def save_cursor(c):
    with open(STATE, "w") as f:
        f.write(c)

def init_db(db):
    db.execute("""
    CREATE TABLE IF NOT EXISTS events (
        ts INTEGER,
        ip TEXT,
        port INTEGER,
        user TEXT,
        type TEXT
    )
    """)

def classify(line):
    if not line or not line[0].isdigit():
        return None

    try:
        ts_str, msg = line.split(" ", 1)
        ts = int(float(ts_str))
    except:
        return None

    m = re.search(r"Invalid user ([a-zA-Z0-9_-]+) from ([0-9a-fA-F:.]+) port (\d+)", msg)
    if m:
        return ts, m.group(2), int(m.group(3)), m.group(1), "invalid_user"
    m = re.search(r"PAM: Authentication failure.*from ([0-9a-fA-F:.]+)", msg)
    if m:
        return ts, m.group(1), None, None, "pam_failure"
    m = re.search(r"Received disconnect from ([0-9a-fA-F:.]+) port (\d+)", msg)
    if m:
        return ts, m.group(1), int(m.group(2)), None, "disconnect"
    m = re.search(r"Connection closed by .* from ([0-9a-fA-F:.]+) port (\d+)", msg)
    if m:
        return ts, m.group(1), int(m.group(2)), None, "connection_closed"
    m = re.search(r"Failed .* for (?:invalid user )?([a-zA-Z0-9_-]+) from ([0-9a-fA-F:.]+) port (\d+)", msg)
    if m:
        return ts, m.group(2), int(m.group(3)), m.group(1), "failed_auth"
    return None

def geo(db, ip):
    row = db.execute("SELECT country, lat, lon FROM ip_geo WHERE ip=?", (ip,)).fetchone()
    if row:
        return row
    try:
        r = requests.get(f"http://ip-api.com/json/{ip}", timeout=5).json()
        if r.get("status") == "success":
            data = (r["country"], r["lat"], r["lon"])
        else:
            data = (None, None, None)
    except:
        data = (None, None, None)
    db.execute(
        "INSERT OR REPLACE INTO ip_geo VALUES (?, ?, ?, ?)",
        (ip, *data)
    )
    time.sleep(0.1)
    print(f"Get location for {ip} to {data[0]}")
    return data

def main():
    db = sqlite3.connect(DB)

    init_db(db)

    db.execute("""
    CREATE TABLE IF NOT EXISTS ip_geo (
        ip TEXT PRIMARY KEY,
        country TEXT,
        lat REAL,
        lon REAL
    )
    """)

    cmd = ["journalctl", "-u", "sshd", "-o", "short-unix", "--no-pager"]

    cur = get_cursor()
    if cur:
        cmd += ["--after-cursor", cur]

    out = subprocess.check_output(cmd, text=True)
    lines = out.splitlines()

    added = 0
    for line in lines:
        result = classify(line)
        if not result:
            continue

        ts, ip, port, user, typ = result
        db.execute(
            "INSERT INTO events VALUES (?, ?, ?, ?, ?)",
            (ts, ip, port, user, typ)
        )
        added += 1
        if ip:
            geo(db, ip)


    db.commit()

    # save cursor
    out = subprocess.check_output(
        ["journalctl", "-u", "sshd", "-n", "1", "-o", "export"],
        text=True
    )

    for l in out.splitlines():
        if l.startswith("__CURSOR="):
            save_cursor(l.split("=", 1)[1])
            break

    db.close()

    print(f"Parsed {len(lines)} lines and added {added} events")

if __name__ == "__main__":
    main()