import psutil
import flask
import argparse
import json
import datetime

app = flask.Flask("ssh-finder")

def makeResponseFromProcess(p):

    lastSU = None
    lastZSH = None

    for subProc in p.children(recursive=True):
        if "zsh" in subProc.name():
            lastZSH = subProc
        if "su" in subProc.name():
            lastSU = subProc

    print(lastZSH)
    # check cwd #
    if lastZSH == None:
        return None
    cwd = lastZSH.cwd()
    print("cwd", cwd)

    # check user if exits #
    user = None
    if lastSU:
        user = lastSU.cmdline()[1] 
    if user and not user.isalnum():
        raise AssertionError("User not alphanum: {}".format(user))

    return { "cwd" : cwd, "user" : user }

def findSSH(createTime, graceTime):
    for p in psutil.process_iter():
        try:
            if 'ssh' in p.name().lower():
                for p in p.children():
                    if 'ssh' in p.name().lower():
                        dtProcess = datetime.datetime.fromtimestamp(p.create_time())
                        if dtProcess - createTime < graceTime:
                            print(p)
                            response = makeResponseFromProcess(p)
                            if response:
                                return response
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass

@app.route('/find')
def find():
    time = int(float(flask.request.args.get("time")))
    #grace = datetime.timedelta(seconds=5)
    grace = datetime.timedelta(hours=500)
    dt = datetime.datetime.fromtimestamp(time)
    data = findSSH(dt, grace)
    if data:
        r = flask.Response(json.dumps(data), 200)
        r.headers = { "Content-Type" : "application/json" }
        return r
    else:
        return ("", 404)

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='SSH Session Finder')
    parser.add_argument('--interface', default="localhost", 
            help='Interface on which flask (this server) will take requests on')
    parser.add_argument('--port', default="5000",
            help='Port on which flask (this server) will take requests on')
   
    args = parser.parse_args()
    app.run(host=args.interface, port=args.port)