# main logic 

from sys_init import setup
setup()

import pyinotify
import subprocess
import re
import pwd
import os
import json
import atexit
import fcntl

def get_ssh_port(pid):
    '''
    /var/log/secure
    Jan 31 07:50:28 vultr sshd[43690]: Accepted publickey for root from 210.10.76.5 port 43730 ssh2: ED25519 SHA256:qz9ffMCb3vPlabn3ZHee00qIPBxkDiUiVSorcUkGdII
    Jan 31 07:50:28 vultr sshd[43690]: pam_unix(sshd:session): session opened for user root(uid=0) by root(uid=0)
    Jan 31 07:50:29 vultr sshd[43693]: Received disconnect from 210.10.76.5 port 43730:11: disconnected by user
    Jan 31 07:50:29 vultr sshd[43693]: Disconnected from user root 210.10.76.5 port 43730
    Jan 31 07:50:29 vultr sshd[43690]: pam_unix(sshd:session): session closed for user root
    '''
    '''
    lsof -i -n
    sshd      1845   root    4u  IPv4  23137      0t0  TCP 45.32.108.159:ssh->210.10.76.5:45460 (ESTABLISHED)
    sshd      1848   root    4u  IPv4  23137      0t0  TCP 45.32.108.159:ssh->210.10.76.5:45460 (ESTABLISHED)
    sshd      1848   root    8u  IPv4  23259      0t0  TCP *:44699 (LISTEN)
    sshd      1848   root    9u  IPv6  23260      0t0  TCP *:44699 (LISTEN)
    '''
    '''
    # less efficient but readable
    pid = '33216'
    pids = []
    port = -1
    for l in lines:
        l = l.split()
        if pid == l[1]:
            commonstring = l[8]
            for l2 in lines:
                if commonstring in l2.split():
                    pids.append(l2.split()[1])
    for i in pids:
        for l in lines:
            l = l.split()
            if i in l[1] and "*:" in l[8]:
                port = l[8][2:]
    '''
    # using the pid of the ssh connection, get the string of the connection 45.32.108.159:ssh->210.10.76.5:45460 
    # then get the other pid with the same string
    # then get the other pid
    # then get the publicly forwarded port
    try:
        lines = subprocess.check_output("lsof -i -n | grep sshd", shell=True, text=True).splitlines()
        pid = str(pid)
        pids = {l.split()[1] for l in lines
                if pid == l.split()[1] or
                l.split()[8] in [x.split()[8] for x in lines if x.split()[1] == pid]}

        port = next(l.split()[8][2:] for l in lines if l.split()[1] in pids and "*:" in l.split()[8])
        return port
    except:
        return -1

def get_keyname(fingerprint):
    auth_file = os.path.join(os.path.expanduser(f"~{user}"), ".ssh", "authorized_keys")
    try:
        result = subprocess.run(['ssh-keygen', '-lf', f'{auth_file}'], capture_output=True, text=True, shell=False, check=True)   
        # Raises CalledProcessError if command fails, for check = true

        #256 SHA256:d9CBxSLBLzLjYGZDCsO6V+lddlN/elK1hGDBS56cbTo user@host (ED25519)        
        for line in result.stdout.strip().split('\n'):
            parts = line.split()
            if fingerprint in parts[1]:
                return parts[2]
        return 'not found'
        
    except subprocess.CalledProcessError as e:
        print(f"Error executing command: {e}")
        return None
    except Exception as e:
        print(f"Error processing output: {e}")
        return None

def write_data(data):
    with open('/tmp/ssh_sessions.json', 'w') as f:
        # Get exclusive lock for writing
        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
        try:
            # Clear file and write new data
            f.seek(0)
            json.dump(data, f)
            f.truncate()
        finally:
            fcntl.flock(f.fileno(), fcntl.LOCK_UN)

def handle_log_change(event):
    global last_position
    with open(event.pathname, 'r') as f:
        # Go to end and read new lines
        f.seek(last_position)
        new_lines = f.readlines()
        last_position = f.tell()
        
        for line in new_lines:
            #print(f"New log: {line.strip()}")
            line = line.strip() 
            if f"Accepted publickey for {user}" in line:
                pid = re.search(r'\[(\d+)\]', line.split()[4]).group(1)
                port = get_ssh_port(pid)
                keyname = get_keyname(line.split()[15]) 
                srcip = line.split()[10]
                #print(pid, port, keyname, srcip)
                #ssh_sessions[pid] = [srcip, keyname, port]

                ssh_sessions[srcip] = {
                    'srcip': srcip,
                    'pid' : pid,
                    'key': keyname,
                    'pubport': port,
                }
                write_data(ssh_sessions)
            if "pam_unix(sshd:session): session closed" in line:
                pid = re.search(r'\[(\d+)\]', line.split()[4]).group(1)
                if pid in ssh_sessions:
                    del ssh_sessions[pid]
                else:
                    print("WARNING! PID NOT FOUND IN SESSION CACHE!")
                write_data(ssh_sessions) 

            print("db: ",ssh_sessions)



user = "root"

# init last_position
global last_position
with open('/var/log/secure', 'r') as f:
    f.seek(0, 2)
    last_position = f.tell()

global ssh_sessions 
ssh_sessions = {}

wm = pyinotify.WatchManager()
wm.add_watch('/var/log/secure', pyinotify.IN_MODIFY, handle_log_change)
notifier = pyinotify.Notifier(wm)

print("Watching /var/log/secure...")
notifier.loop()