import json
import sys
import random
import os
import argparse
import hashlib

def ccgen(seed=""): # Color code generator
  if(not os.isatty(1)): # Do not output color code at all if output is not a TTY that can accept them
    return ""

  match(seed): # Match color keywords to control codes
    case "RED":
      return "\x1b[31m"
    case "GREEN":
      return "\x1b[32m"
    case "YELLOW":
      return "\x1b[33m"
    case "BLUE":
      return "\x1b[34m"
    case "PURPLE":
      return "\x1b[35m"
    case "CYAN":
      return "\x1b[36m"
    case "GREY":
      return "\x1b[37m"
    case "RESET":
      return "\x1b[0m"

  if(seed): # If seed is provided, use it to generate a color
    random.seed(seed)
    return "\x1b[38;2;{r};{g};{b}m".format(r=random.randrange(50, 256), g=random.randrange(50, 256), b=random.randrange(50, 256))
  return "\x1b[0m"

def ttylog(mystring): #mystring
  if(os.isatty(1)):
    print(mystring)

def ppip(addr, spacing = False, ttycolor = False): # Pretty-print an address with colour and, if it's IPv4 or hashed, a consistent width
  rs = addr
  if(ttycolor):
    rs = ccgen(rs) + rs + ccgen()
  if(len(addr) <= 16 and spacing):
    rs = rs + (" " * (15 - len(addr)))
  return rs

def hashgen(seed, length=16):
  if(length % 2):
    raise Warning("Due to strange behaviour on the part of the hexdigest() function, length should be a multiple of 2")
  return hashlib.shake_256(str(seed).encode("utf-8")).hexdigest(int(length / 2))

cmdping_commands = ['uname', 'nproc', 'echo'] # List of commands that, if used exclusively, qualify a session as a "command ping"

# ArgParse crap coming in hot

parser = argparse.ArgumentParser(
  description="Script to parse JSON logs from popular SSH honeypot Cowrie and print them in a more readable format"
)

parser.add_argument('filepath')

parser.add_argument('-v', '--verbose', action='store_true', help="Make the output more verbose (print abridged logs for all logins where commands are run, not just the interesting ones)")
parser.add_argument('-V', '--verboser', action='store_true', help="Make the output even more verbose (print timestamp, event id, and message (if applicable) for all events)")
parser.add_argument('-W', '--verbosest', action='store_true', help="Make the output even more verboserer (just group the input log by session id and dump it to pretty-printed JSON)")

parser.add_argument('-P', '--pings', action='store_true', help='Include only connections where no authentication attempt was made')
parser.add_argument('-F', '--fails', action='store_true', help='Include only failed logins')
parser.add_argument('-L', '--logins', action='store_true', help='Include only successful logins where no commands were run')
parser.add_argument('-C', '--commands', action='store_true', help='Include only logins where only this list of basic test commands were run: ' + str(cmdping_commands))
parser.add_argument('-E', '--eventful', action='store_true', help='Include only the cool interesting sessions that aren\'t any of that other stuff')

parser.add_argument('-r', '--redact-ip', action='store_true', help='Redact IP addresses by converting them to 16-character hashes')
parser.add_argument('-i', '--filter-ip', help='Filter data to a specific IP or censored hash thereof')

parser.add_argument('-s', '--stats', action='store_true', help='Print aggregated statistics')
parser.add_argument('-S', '--only-stats', action='store_true', help='Print only aggregated statistics')


args = parser.parse_args()
print_stats = args.stats or args.only_stats

verbosity = 0
if(args.verbosest):
  verbosity = 3
elif(args.verboser):
  verbosity = 2
elif(args.verbose):
  verbosity = 1

filepath = args.filepath

blankmode = (args.pings or args.fails or args.logins or args.commands or args.eventful or args.only_stats)

include_empty_pings   = (not blankmode or args.pings)
include_failed_logins = (not blankmode or args.fails)
include_login_pings   = (not blankmode or args.logins)
include_command_pings = (not blankmode or args.commands)
include_eventful      = (not blankmode or args.eventful)

sessions = {} # Dict - session_id : list of events in session

def session_ip(session): # Grab IP from session
  return session[0]['src_ip']

def session_timestamp(session): # Grab initial timestamp of session
  return session[0]['timestamp']

def session_duration(session): # Grab duration of session -- if session never closed return -1
  if(session[-1]['eventid'] == "cowrie.session.closed"):
    return float(session[-1]['duration'])
  return -1

def session_sensor(session): # Grab sensor from session
  return session[0]['sensor']

def session_attempted_login(session): # Determine whether a login was attempted
  for event in session:
    if(event['eventid'] == "cowrie.login.failed" or event['eventid'] == "cowrie.login.success"):
      return True
  return False

def session_login_succeeded(session): # Determine whether a login attempt was successful
  for event in session:
    if(event['eventid'] == "cowrie.login.success"):
      return True
  return False

def session_cred_combos(session): # List passwords used in login attempt
  ra = []
  if(session_attempted_login(session)):
    for event in session:
      if(event['eventid'] == "cowrie.login.failed" or event['eventid'] == "cowrie.login.success"):
        if('password' in event):
          ra.append((event['username'], event['password']))
        elif('key' in event):
          ra.append((event['username'], []))
  return ra

def session_ran_command(session): # Was a command run during the session
  for event in session:
    if(event['eventid'] == "cowrie.command.input"):
      return True
  return False

def session_saved_file(session):
  for event in session:
    if(event['eventid'] == "cowrie.session.file_download"):
      return True
  return False

def session_is_cmdping(session): # Determine if a session qualifies as a "command ping"
  if(session_ran_command(session)):
    for event in session:
      if(event['eventid'] == "cowrie.command.input"):
        argv0 = event['input'] # Isolate the name of the command
        if(" " in argv0):
          argv0 = event['input'][:event['input'].index(" ")]

        if(argv0 not in cmdping_commands):
          return False
    return True
  return False

def session_pingcmds(session): # Return array of all "command ping" commands used in a session
  ra = []
  if(session_ran_command(session)):
    for event in session:
      if(event['eventid'] == "cowrie.command.input"):
        argv0 = event['input'] # Isolate the name of the command
        if(" " in argv0):
          argv0 = event['input'][:event['input'].index(" ")]

        if(argv0 in cmdping_commands):
          ra.append(argv0)
  return ra

def load_json_to_array(fp):
  with open(fp, "r") as logfile: # Load da file and populate da master array
    for line in logfile:
      entry = json.loads(line)
      
      if(
        (not args.filter_ip or entry['src_ip'] == args.filter_ip or hashgen(entry['src_ip']) == args.filter_ip) # If it's not the IP we're filtering for don't even bother man
        ):
        session_id = entry['session']

        if(args.redact_ip):
          entry['src_ip'] = hashgen(entry['src_ip'])

        if(session_id not in sessions):
          sessions.update({session_id: [entry]})
        else:
          sessions[session_id].append(entry)

if(os.path.isfile(filepath)):
  load_json_to_array(filepath)
elif(os.path.isdir(filepath)):
  # run load_json_to_array for every json file in the directory
  files = {}
  for f in os.listdir(filepath):
    if(f[:11] == "cowrie.json" or f[-5:] == ".json"):
      fi = os.path.join(filepath, f)
      files.update({fi: os.path.getmtime(fi)})
  for g in dict(sorted(files.items(), key=lambda item: item[1])):
    load_json_to_array(g)
else:
  raise FileNotFoundError

# Filtering

newsessions = sessions.copy()

for session_id in sessions:
  session = sessions[session_id]
  
  if(blankmode and (
       (not include_empty_pings   and not session_attempted_login(session))
    or (not include_failed_logins and     session_attempted_login(session) and not session_login_succeeded(session))
    or (not include_login_pings   and     session_login_succeeded(session) and not session_ran_command(session))
    or (not include_command_pings and     session_is_cmdping(session))
    or (not include_eventful      and     session_ran_command(session)     and not session_is_cmdping(session))
  )):
    newsessions.pop(session_id)

sessions = newsessions

# Printing logic coming up

if(verbosity == 2): # High verbosity mode: print timestamp, event type and message (if applicable) for every event. I could probably do better but I'm lazy
  for session_id in sessions:
    session = sessions[session_id]
    ip = session_ip(session_id)
    print("\n--- Session " + ccgen(session_id) + session_id + ccgen() + " from " + ppip(ip, ttycolor=True) + " to " + session_sensor(session) + " ---")
    for entry in session:
      eventid = entry['eventid']
      if(entry['eventid'] == "cowrie.session.params"):
        print(entry["timestamp"] + " > " + eventid + (" " * (29 - len(eventid))) + ": ARCH: " + entry["arch"])
      else:
        print(entry["timestamp"] + " > " + eventid + (" " * (29 - len(eventid))) + ": " + entry['message'])

elif(verbosity == 3): # Maximum verbosity: literally just dump the now-sorted-by-session logs as pretty-printed JSON
  print(json.dumps(sessions, indent=4))

else: # Default (low) verbosity: just the good stuff
  for session_id in sessions:
    
    # Declaring initial variables
    session = sessions[session_id]
    ip = session_ip(session)
    timestamp = session_timestamp(session)[:22].replace("T", " ") # clean up the timestamp text all pretty like
    durationstring = "closed after " + str(session_duration(session)) + "s"
    if(session_duration(session) < 0):
      durationstring = "active as of processing"

    if(not session_attempted_login(session)): # Printing empty pings
      print(timestamp + " > " + ppip(ip, True, True) + " > Connection opened and " + durationstring)

    elif(session_attempted_login(session) and not session_login_succeeded(session)): # Printing failed logins
      print(timestamp + " > " + ppip(ip, True, True) + " > " + ccgen("RED") +"Failed login" + ccgen() + " with credentials: " + str(session_cred_combos(session)))

    elif(session_ran_command(session)): # Printing logic for sessions in which a command was run

      if(not verbosity == 1 and session_is_cmdping(session)): # If -v is not enabled, collapse "command pings" to one-liners
        print(timestamp + " > " + ppip(ip, True, True) + " > " + ccgen("GREEN") + "Login" + ccgen() + " with credentials: " + str(session_cred_combos(session)) + ccgen("BLUE") + " ran commands " + ccgen() + str(session_pingcmds(session)))

      else: # Printing logic for un-collapsed sessions where commands were run
        print(timestamp + " > " + ppip(ip, True, True) + " > " + ccgen("CYAN") + "Session" + ccgen() + ":")
        for entry in session:
          match(entry["eventid"]):
            case "cowrie.login.failed":
              pw = []
              if("password" in entry):
                pw = entry["password"]
              print(" | Failed login: " + entry["username"] + "/" + str(pw))
            case "cowrie.login.success":
              print(" | Successful login: " + entry["username"] + "/" + entry["password"])
            case "cowrie.command.input":
              print(" | CMD: " + entry["input"])
            case "cowrie.session.file_download":
              df = ""
              if("url" in entry):
                df = entry["url"]
              elif("destfile" in entry):
                df = entry["destfile"]
              if(verbosity == 1):
                print(" | Saved file " + df + " to " + entry["outfile"])
              else:
                print(" | Saved file " + df + " to " + entry["outfile"][:40] + "...")
        print(" X Session " + durationstring)
        
    else: # Printing logic for successful logins in which no commands were run
      print(timestamp + " > " + ppip(ip, True, True) + " > " + ccgen("GREEN") + "Login" + ccgen() + " with credentials: " + str(session_cred_combos(session)))

if(args.stats):
  print("Aggregate statistics TODO")