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

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

PATTERN = re.compile(r"PAM: Authentication failure for root from ([0-9a-fA-F:.]+)")

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

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

def geo(ip):
    print(f"Send geo request for {ip}")
    r = requests.get(f"http://ip-api.com/json/{ip}", timeout=5).json()
    if r["status"] == "success":
        return r["country"], r["lat"], r["lon"]
    return None, None, None

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

    db.execute("CREATE TABLE IF NOT EXISTS failures (ts INTEGER, ip TEXT)")

    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()

    print(f"Must parse {len(lines)}")
    for line in lines:
        print(line)
        if not line or not line[0].isdigit():
            continue

        ts_str, msg = line.split(" ", 1)

        try:
            ts = int(float(ts_str))
        except ValueError:
            continue

        m = PATTERN.search(msg)
        if not m:
            continue

        ip = m.group(1)

        # store raw event
        print(f"Inserting {ip}")
        db.execute("INSERT INTO failures VALUES (?, ?)", (ts, ip))

        # check cache
        cached = db.execute(
            "SELECT 1 FROM ip_geo WHERE ip = ?",
            (ip,)
        ).fetchone()

        if not cached:
            country, lat, lon = geo(ip)
            db.execute(
                "INSERT OR REPLACE INTO ip_geo VALUES (?, ?, ?, ?)",
                (ip, country, lat, lon)
            )
            time.sleep(0.2)

    db.commit()
    db.close()

    # 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

if __name__ == "__main__":
    main()
    print("Parsed everything")