#!python

"""
 Copyright (c) 2019 Cisco Systems, Inc. All rights reserved.
 License at https://github.com/cisco/mercury/blob/master/LICENSE
"""

import os
import sys
import json
import pcap
import time
import optparse
from socket import AF_INET, AF_INET6, inet_ntop

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.dirname(os.path.abspath(__file__))+'/../')
from pmercury.utils.config_parser import parse_config
from pmercury.utils.packet_proc import pkt_proc

from pmercury.protocols.tcp import TCP
from pmercury.protocols.tls import TLS
from pmercury.protocols.dtls import DTLS
from pmercury.protocols.http import HTTP
from pmercury.protocols.dhcp import DHCP
from pmercury.protocols.iquic import IQUIC
from pmercury.protocols.tls_server import TLS_Server
from pmercury.protocols.dtls_server import DTLS_Server
from pmercury.protocols.http_server import HTTP_Server
from pmercury.protocols.tls_certificate import TLS_Certificate





class Fingerprinter:

    def __init__(self, database, output=None, analyze=False, num_procs=0,
                 human_readable=False, group=False, endpoint=False, experimental=False, sslkeylog=None):

        self.analyze        = analyze
        self.num_procs      = num_procs
        self.human_readable = human_readable
        self.group          = group
        self.endpoint       = endpoint
        self.FLOW_TIMEOUT   = 150
        self.FLOW_UPDATE    = 15

        self.fast_path = True
        if analyze or human_readable or group or endpoint or experimental:
            self.fast_path = False

        if output == None:
            self.out_file_pointer = sys.stdout
        else:
            self.out_file_pointer = open(output, 'w', buffering=8192*64)

        if endpoint:
            if output == None:
                self.endpoint_file_pointer = sys.stdout
            else:
                self.endpoint_file_pointer = open('endpoint_'+output, 'w', buffering=8192*64)

        config = parse_config('../config.yaml')

        # register parsers
        self.parsers = {}
        self.parsers['tcp']          = TCP('resources/fingerprint_db_tcp.json.gz')
        self.parsers['tls']          = TLS(database)
        self.parsers['tls_server']   = TLS_Server()
        self.parsers['dtls']         = DTLS(database)
        self.parsers['dtls_server']  = DTLS_Server()
        self.parsers['http']         = HTTP(None, config)
        self.parsers['http_server']  = HTTP_Server(None, config)
        self.parsers['server_certs'] = TLS_Certificate()
        self.parsers['dhcp']         = DHCP(None, config)
        self.parsers['iquic']        = IQUIC()

        self.contextual_data = ['tls','http','http_server','dhcp']

        self.flow_parsers  = []
        self.experimental = False
        if experimental == True:
            self.experimental = True
            from pmercury.protocols.ssh import SSH
            from pmercury.protocols.tls_decrypt import TLS_Decrypt

            self.parsers['ssh'] = SSH()
            self.flow_parsers.extend([('ssh', SSH())])
            if sslkeylog != None:
                self.parsers['tls_decrypt'] = TLS_Decrypt(sslkeylog)
                self.flow_parsers.append(('tls_decrypt', TLS_Decrypt(sslkeylog)))

        if endpoint == True:
            from pmercury.utils.endpoint import Endpoints
            self.endpoint_model = Endpoints()

        # data
        self.flow_cache = {}


    def process(self, input_files):
        if len(input_files) == 0:
            print('error: need a pcap/interface')
            return 1

        for input_file in input_files:
            if '.pcap' in input_file:
                if not os.path.exists(input_file):
                    print('error: file does not exist')
                    return 1
                self.process_pcap(input_file)
            else:
                self.process_capture(input_file)


    def process_pcap(self, pcap_file):
        p = pcap.pcap(pcap_file)
#        p.setfilter('ip proto 6 or ip6 proto 6 or ip proto 17 or ip6 proto 17')
        p.dispatch(-1, self.process_packet)

        if self.group:
            self.write_flows(active=False)
        if self.endpoint:
            self.endpoint_model.write_all(self.endpoint_file_pointer)


    def process_capture(self, iface):
        import signal
        def signal_handler(signal, frame):
            if self.group:
                self.write_flows(active=False)
            sys.exit(0)
        signal.signal(signal.SIGINT, signal_handler)

        p = pcap.pcap(iface)
#        p.setfilter('ip proto 6 or ip6 proto 6 or ip proto 17 or ip6 proto 17')
        flow_timer = time.time()
        while 1:
            p.dispatch(-1, self.process_packet)

            if self.group and time.time()-flow_timer > self.FLOW_UPDATE:
                flow_timer = time.time()
                self.write_flows(active=True)

            nrecv, ndrop, nifdrop = p.stats()
            print('nrecv: % 12i, ndrop: % 12i, nifdrop: % 12i' % (nrecv, ndrop, nifdrop))
        if self.group:
            self.write_flows(active=False)


    def process_packet(self, ts, buf):
        flow = pkt_proc(ts, buf)

        if flow != None:
            if self.fast_path:
                self.write_record(flow)
                return

            if flow['fingerprints'] != {}:
                fp_type = next(iter(flow['fingerprints']))
                if self.human_readable:
                    desc_ = self.parsers[fp_type].get_human_readable(flow['fingerprints'][fp_type])
                    if desc_ != None:
                        if 'readable' not in flow:
                            flow['readable'] = {}
                        flow['readable'][fp_type] = desc_
                    if fp_type == 'tls_server':
                        if 'tls' in flow and 'server_certs' in flow['tls']:
                            desc_ = self.parsers['server_certs'].get_human_readable(flow['tls']['server_certs'])
                            if desc_ != None:
                                if 'readable' not in flow:
                                    flow['readable'] = {}
                                flow['readable']['server_certs'] = desc_
                if self.analyze:
                    context = None
                    if fp_type in flow:
                        context = flow[fp_type]
                    proc_info = self.parsers[fp_type].proc_identify(flow['fingerprints'][fp_type], context,
                                                                    flow['dst_ip'], flow['dst_port'], self.num_procs)
                    if proc_info != None:
                        flow['analysis'] = proc_info
            elif 'tls' in flow and 'server_certs' in flow['tls']:
                desc_ = self.parsers['server_certs'].get_human_readable(flow['tls']['server_certs'])
                if desc_ != None:
                    if 'readable' not in flow:
                        flow['readable'] = {}
                    flow['readable']['server_certs'] = desc_

            if self.group:
                self.update_flow_cache(flow)
            else:
                self.write_record(flow)

            if self.endpoint:
                self.endpoint_model.update(flow)

        if self.experimental:
            self.process_packet_flow(ts, buf)


    def process_packet_flow(self, ts, buf):
        if buf[12] == 0x08 and buf[13] == 0x00: # IPv4
            ip_type = 'ipv4'
            ip_length = 20
            ip_offset = 14
            protocol = buf[23]
        elif buf[12] == 0x86 and buf[13] == 0xdd: # IPv6
            ip_type = 'ipv6'
            ip_length = 40
            ip_offset = 14
            protocol = buf[20]
        elif buf[14] == 0x08 and buf[15] == 0x00: # IPv4 (hack for linux cooked capture)
            ip_type = 'ipv4'
            ip_length = 20
            ip_offset = 16
            protocol = buf[25]
        else: # currently skip other types
            return -1

        data_len = len(buf)
        if protocol == 6:
            tcp_offset = ip_offset+ip_length
            if tcp_offset+20 > data_len:
                return -1
            tcp_length = (buf[tcp_offset+12] >> 0x04)*4
            app_offset = tcp_offset + tcp_length

            if self.experimental:
                for _,fparser_ in self.flow_parsers:
                    ffp_type, ffp_str_, fcontext_ = fparser_.fingerprint(buf, ip_offset, tcp_offset, app_offset,
                                                                         ip_type, ip_length, data_len)
                    if ffp_str_ != None:
                        self.process_result(buf, ffp_type, ffp_str_, fcontext_, tcp_offset,
                                            ip_type, ip_offset, ip_length, protocol, ts)


    def process_result(self, buf, fp_type, fp_str_, context_, tcp_offset, ip_type, ip_offset, ip_length, protocol, ts):
        src_port = int.from_bytes(buf[tcp_offset:tcp_offset+2], byteorder='big')
        dst_port = int.from_bytes(buf[tcp_offset+2:tcp_offset+4], byteorder='big')
        if ip_type == 'ipv4':
            o_ = ip_offset+ip_length-8
            src_ip = inet_ntop(AF_INET, buf[o_:o_+4])
            o_ += 4
            dst_ip = inet_ntop(AF_INET, buf[o_:o_+4])
        else:
            o_ = ip_offset+ip_length-32
            src_ip = inet_ntop(AF_INET6, buf[o_:o_+16])
            o_ += 16
            dst_ip = inet_ntop(AF_INET6, buf[o_:o_+16])

        flow = {'src_ip':src_ip,
                'dst_ip':dst_ip,
                'src_port':src_port,
                'dst_port':dst_port,
                'protocol':protocol,
                'event_start':ts,
                'fingerprints':{fp_type: fp_str_}}
        if context_ != None:
            for x_ in context_:
                if x_['data'] != None:
                    if fp_type not in flow:
                        flow[fp_type] = {}
                    flow[fp_type][x_['name']]  = x_['data']
        if self.human_readable:
            hr_fp_type = fp_type
            if 'tls_decrypt' in hr_fp_type:
                hr_fp_type = 'tls_decrypt'
            desc_ = self.parsers[hr_fp_type].get_human_readable(fp_str_)
            if desc_ != None:
                if 'readable' not in flow:
                    flow['readable'] = {}
                flow['readable'][fp_type] = desc_
        if self.analyze:
            if fp_type in self.parsers:
                proc_info = self.parsers[fp_type].proc_identify(fp_str_, context_, dst_ip, dst_port, self.num_procs)
                if proc_info != None:
                    flow['analysis'] = proc_info

        if self.group:
            self.update_flow_cache(flow)
        else:
            self.write_record(flow)


    def process_csv(self, fp_type, fp_str, src_ip, dst_ip, src_port, dst_port, protocol, ts, context, debug=None):
        flow = {'src_ip':src_ip,
                'dst_ip':dst_ip,
                'src_port':src_port,
                'dst_port':dst_port,
                'protocol':protocol,
                'event_start':ts,
                'fingerprints':{fp_type: fp_str}}

        if context != None:
            flow[fp_type] = {}
            for k in context:
                flow[fp_type][k] = context[k]

        if self.human_readable:
            desc_ = self.parsers[fp_type].get_human_readable(flow['fingerprints'][fp_type])
            if desc_ != None:
                if 'readable' not in flow:
                    flow['readable'] = {}
                flow['readable'][fp_type] = desc_

        if self.analyze:
            endpoint = None
            if self.endpoint:
                endpoint = self.endpoint_model.get_endpoint(src_ip)
            proc_info = self.parsers[fp_type].proc_identify(flow['fingerprints'][fp_type], context, flow['dst_ip'],
                                                            flow['dst_port'], self.num_procs, endpoint, False, debug)
            if proc_info != None:
                flow['analysis'] = proc_info

            os_info = self.parsers[fp_type].os_identify(flow['fingerprints'][fp_type], self.num_procs)
            if os_info != None:
                if 'analysis' not in flow:
                    flow['analysis'] = {}
                flow['analysis']['os_info'] = os_info

        if self.endpoint:
            self.endpoint_model.update(flow)

        return flow


    def update_flow_cache(self, flow):
        t_flow_key, t_r_flow_key = self.get_flow_keys(flow)
        if t_flow_key in self.flow_cache:
            flow_key = t_flow_key
        elif t_r_flow_key in self.flow_cache:
            flow_key = t_r_flow_key
        else:
            flow_key = t_flow_key
            self.init_flow_cache(flow, flow_key)

        # update end time
        self.flow_cache[flow_key]['event_end'] = flow['event_start']

        # add fingerprint
        for k in flow['fingerprints']:
            packet_fp = {}
            packet_fp['event_start'] = flow['event_start']
            packet_fp[k] = flow['fingerprints'][k]
            if k in flow:
                if k not in self.flow_cache[flow_key]:
                    self.flow_cache[flow_key][k] = {}
                for k1 in flow[k]:
                    self.flow_cache[flow_key][k][k1] = flow[k][k1]
            if 'analysis' in flow:
                packet_fp['analysis'] = flow['analysis']
            self.flow_cache[flow_key]['fingerprints'].append(packet_fp)

        # add contextual data
        for k in self.contextual_data:
            if k in flow:
                if k not in self.flow_cache[flow_key]:
                    self.flow_cache[flow_key][k] = {}
                for k1 in flow[k]:
                    self.flow_cache[flow_key][k][k1] = flow[k][k1]

        if 'readable' in flow:
            if 'readable' not in self.flow_cache[flow_key]:
                self.flow_cache[flow_key]['readable'] = {}
            for k in flow['readable']:
                self.flow_cache[flow_key]['readable'][k] = flow['readable'][k]




    def init_flow_cache(self, flow, flow_key):
        self.flow_cache[flow_key] = {}
        self.flow_cache[flow_key]['src_ip']       = flow['src_ip']
        self.flow_cache[flow_key]['dst_ip']       = flow['dst_ip']
        self.flow_cache[flow_key]['src_port']     = flow['src_port']
        self.flow_cache[flow_key]['dst_port']     = flow['dst_port']
        self.flow_cache[flow_key]['protocol']     = flow['protocol']
        self.flow_cache[flow_key]['event_start']  = flow['event_start']
        self.flow_cache[flow_key]['event_end']    = flow['event_start']
        self.flow_cache[flow_key]['fingerprints'] = []


    def get_flow_keys(self, flow):
        src_ip   = flow['src_ip']
        dst_ip   = flow['dst_ip']
        src_port = str(flow['src_port'])
        dst_port = str(flow['dst_port'])
        protocol = str(flow['protocol'])

        fk   = src_ip + ':' + dst_ip + ':' + src_port + ':' + dst_port + ':' + protocol
        r_fk = dst_ip + ':' + src_ip + ':' + dst_port + ':' + src_port + ':' + protocol

        return fk, r_fk


    def lookup_fingerprint_string(self, fp_str):
        for protocol_type in self.parsers:
            if not hasattr(self.parsers[protocol_type], 'get_database_entry'):
                continue
            fp_ = self.parsers[protocol_type].get_database_entry(fp_str, None)
            if fp_ != None:
                self.write_record(fp_)


    def get_database_entry(self, fp_str, fp_type):
        if not hasattr(self.parsers[fp_type], 'get_database_entry'):
            return None
        return self.parsers[fp_type].get_database_entry(fp_str, None)


    def write_record(self, flow_repr):
        if 'fingerprints' in flow_repr and flow_repr['fingerprints'] == {}:
            del flow_repr['fingerprints']

        try:
            self.out_file_pointer.write('%s\n' % json.dumps(flow_repr))
        except:
            print('error: flow incorrectly formatted:')
            print(flow_repr)
            print()


    def write_flows(self, active=False):
        curtime = time.time()
        flow_keys = [k for k in self.flow_cache.keys()
                     if not active or
                     ((curtime - self.flow_cache[k]['event_start']) > self.FLOW_TIMEOUT)]

        for k in flow_keys:
            self.flow_cache[k]['event_start'] = self.flow_cache[k]['event_start']
            self.flow_cache[k]['event_end'] = self.flow_cache[k]['event_end']
            self.write_record(self.flow_cache[k])
            del self.flow_cache[k]


    def close_files(self):
        if self.out_file_pointer != None and self.out_file_pointer != sys.stdout:
            self.out_file_pointer.close()
        if self.endpoint and self.endpoint_file_pointer != None:
            self.endpoint_file_pointer.close()


def main():
    start = time.time()

    parser = optparse.OptionParser()

    parser.add_option('-c','--capture',action='store',dest='capture_interface',
                      help='live packet capture',default=None)
    parser.add_option('-r','--read',action='store',dest='pcap_file',
                      help='read packets from file',default=None)
    parser.add_option('-d','--fp_db',action='store',dest='fp_db',
                      help='location of fingerprint database',default='resources/fingerprint_db.json.gz')
    parser.add_option('-f','--fingerprint',action='store',dest='output',
                      help='write fingerprints to file',default=None)
    parser.add_option('-l','--lookup',action='store',dest='lookup',
                      help='lookup fingerprint string <fp_str>',default=None)
    parser.add_option('-a','--analysis',action='store_true',dest='analyze',
                      help='perform process identification',default=False)
    parser.add_option('-w','--human-readable',action='store_true',dest='human_readable',
                      help='return human readable fingerprint information',default=False)
    parser.add_option('-e','--endpoint',action='store_true',dest='endpoint',
                      help='enable endpoint modeling',default=False)
    parser.add_option('-x','--experimental',action='store_true',dest='experimental',
                      help='turns on all experimental features',default=False)
    parser.add_option('-g','--group-flows',action='store_true',dest='group',
                      help='aggregate packet-based fingerprints to flow-based',default=False)
    parser.add_option('-n','--num-procs',action='store',dest='num_procs',type='int',
                      help='return the top-n most probable processes',default=0)
    parser.add_option('-s','--sslkeylogfile',action='store',dest='sslkeylog',
                      help='filename of sslkeylog output for decryption',default=None)

    options, args = parser.parse_args()

    input_files = []
    if options.pcap_file != None:
        input_files.append(options.pcap_file)
    for x in args:
        if '.pcap' in x:
            input_files.append(x)
    if options.capture_interface != None:
        input_files.append(options.capture_interface)

    output = options.output

    fingerprinter = Fingerprinter(options.fp_db, output, options.analyze,
                                  options.num_procs, options.human_readable, options.group,
                                  options.endpoint, options.experimental, options.sslkeylog)

    if options.lookup != None:
        fingerprinter.lookup_fingerprint_string(options.lookup)
    elif len(input_files) > 0:
        fingerprinter.process(input_files)
    else:
        print('error: need a pcap/interface or fingerprint string for lookup')

    fingerprinter.close_files()


if __name__ == '__main__':
    sys.exit(main())
