#!/usr/bin/python

from __future__ import print_function  # logging, python2-only.

"""
A script that reads data generated by p0f -f p0f.log, looking for all entries
about an IP read from NCAT_REMOTE_ADDR environment variable. Then it prints out
all the information it has found. To try it out, run "p0f -i any -o p0f.log"
and ncat -l -k --sh-exec "python p0fme.py".

Script tested under Python versions 2.7 and 3.3.
"""

P0F_LOG_FILE = "p0f.log"

import datetime  # logging
import sys  # logging
import os  # environ
import time  # sleeping to wait for data
import sys  # to flush STDOUT


def expand_ipv6(ip):
    """
    Expands short IPv6 address like ::1 into an expanded form without trailing
    zeros. Copied from:

    http://svn.python.org/projects/python/tags/r31b1/Lib/ipaddr.py

    (Py3 standard library; added some modifications to match p0f's output data)
    """

    new_ip = []
    hextet = ip.split('::')
    sep = len(hextet[0].split(':')) + len(hextet[1].split(':'))
    new_ip = hextet[0].split(':')

    for _ in range(8 - sep):
        new_ip.append('0')
    new_ip += hextet[1].split(':')

    # Now need to make sure every hextet is 4 lower case characters.
    # If a hextet is < 4 characters, we've got missing leading 0's.
    ret_ip = []
    for hextet in new_ip:
        if hextet == '':
                hextet = '0'
        ret_ip.append(hextet.lower())
    return ':'.join(ret_ip)


def split_by_equals(str_):
    ret = str_.split('=')
    return ret[0], ''.join(ret[1:])

if __name__ == "__main__":
    try:
        ip = os.environ['NCAT_REMOTE_ADDR']
        if os.environ['NCAT_PROTO'] != 'TCP':
            sys.exit("ERROR: This script works for TCP servers only!")
    except KeyError:
        sys.exit("ERROR: This script has to be run from inside of Ncat.")
    print("[%s] Got a request from %s" % (
        datetime.datetime.now().isoformat(' '), ip), file=sys.stderr)

    print("Hold on, I'm collecting data on you...")
    sys.stdout.flush()
    time.sleep(3.0)

    if ':' in ip:  # We need to expand IPv6 addresses in a specific way.
            ip = expand_ipv6(ip)
    result = {}

    # Reading the log backward will give us more recent results.
    for line in reversed(open(P0F_LOG_FILE).readlines()):

        without_date = line.split('] ')
        if without_date == ['\n']:
            continue
        without_date = ''.join(without_date[1:])

        # Create a key-value dictionary out of the '|'-separated substrings.
        properties = dict(map(split_by_equals, without_date.split('|')))

        if not properties['cli'].startswith(ip):
            continue  # Not the IP we're looking for, check next one.

        for key in properties:
            if not key in result or result[key] == '???':
                result[key] = properties[key]

    if not result:
        print("Got nothing on you. Try again and I will, though.")

    # Now that we've finished, print out the results.
    for key in sorted(result):
        print("%s: %s" % (key, result[key].rstrip()))