More stats

main.py

@@ -4,25 +4,75 @@ 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)
+    with open(STATE, "w") as f:
+        f.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 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):
+    print(f"Get location for {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)
+    return data
 
 def main():
     db = sqlite3.connect(DB)
 
-    db.execute("CREATE TABLE IF NOT EXISTS failures (ts INTEGER, ip TEXT)")
+    init_db(db)
 
     db.execute("""
     CREATE TABLE IF NOT EXISTS ip_geo (
@@ -34,6 +84,7 @@ def main():
     """)
 
     cmd = ["journalctl", "-u", "sshd", "-o", "short-unix", "--no-pager"]
+
     cur = get_cursor()
     if cur:
         cmd += ["--after-cursor", cur]
@@ -41,45 +92,25 @@ def main():
     out = subprocess.check_output(cmd, text=True)
     lines = out.splitlines()
 
-    print(f"Must parse {len(lines)}")
+    print(f"Parsing {len(lines)} lines")
+
     for line in lines:
-        print(line)
-        if not line or not line[0].isdigit():
+        result = classify(line)
+        if not result:
             continue
 
-        ts_str, msg = line.split(" ", 1)
+        ts, ip, port, user, typ = result
 
-        try:
-            ts = int(float(ts_str))
-        except ValueError:
-            continue
+        print(f"{typ:18} ip={ip} port={port} user={user}")
 
-        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.execute(
+            "INSERT INTO events VALUES (?, ?, ?, ?, ?)",
+            (ts, ip, port, user, typ)
+        )
+        if ip:
+            geo(db, ip)
 
     db.commit()
-    db.close()
 
     # save cursor
     out = subprocess.check_output(
@@ -92,6 +123,9 @@ def main():
             save_cursor(l.split("=", 1)[1])
             break
 
+    db.close()
+    print("Done")
+
 if __name__ == "__main__":
     main()
-    print("Parsed everything")
+    print("All done")