#!/usr/bin/env python
from __future__ import unicode_literals, print_function, division
import signal

__author__ = 'dongliu'

import sys
# check python version
major, minor, = sys.version_info[:2]
if major != 2 or minor < 7:
    print("Python version 2.7.* needed.", file=sys.stderr)
    sys.exit(1)

import argparse
import io

from pcapparser import packet_parser
from pcapparser import pcap, pcapng, utils
from pcapparser.constant import FileFormat
from pcapparser.printer import HttpPrinter
from collections import OrderedDict
import struct

from pcapparser.httpparser import HttpType, HttpParser
from pcapparser import config


# when press Ctrl+C, stop the proxy.
def signal_handler(signal, frame):
    print('\nStopping proxy...')
    sys.exit(0)


signal.signal(signal.SIGINT, signal_handler)


class HttpConn:
    """all data having same source/dest ip/port in one http connection."""
    STATUS_BEGIN = 0
    STATUS_RUNNING = 1
    STATUS_CLOSED = 2
    STATUS_ERROR = -1

    def __init__(self, tcp_pac):
        self.source_ip = tcp_pac.source
        self.source_port = tcp_pac.source_port
        self.dest_ip = tcp_pac.dest
        self.dest_port = tcp_pac.dest_port

        self.status = HttpConn.STATUS_BEGIN

        # start parser thread
        self.processor = HttpPrinter((self.source_ip, self.source_port),
                                     (self.dest_ip, self.dest_port))
        self.http_parser = HttpParser(self.processor)
        self.append(tcp_pac)

    def append(self, tcp_pac):
        if len(tcp_pac.body) == 0:
            return
        if self.status == HttpConn.STATUS_ERROR or self.status == HttpConn.STATUS_CLOSED:
            # not http conn or conn already closed.
            return

        if self.status == HttpConn.STATUS_BEGIN:
            if tcp_pac.body:
                if utils.is_request(tcp_pac.body):
                    self.status = HttpConn.STATUS_RUNNING
        if tcp_pac.source == self.source_ip:
            http_type = HttpType.REQUEST
        else:
            http_type = HttpType.RESPONSE

        if self.status == HttpConn.STATUS_RUNNING and tcp_pac.body:
            self.http_parser.send(http_type, tcp_pac.body)

        if tcp_pac.pac_type == -1:
            # end of connection
            if self.status == HttpConn.STATUS_RUNNING:
                self.status = HttpConn.STATUS_CLOSED
            else:
                self.status = HttpConn.STATUS_ERROR

    def finish(self):
        self.http_parser.finish()


def get_file_format(infile):
    """
    get cap file format by magic num.
    return file format and the first byte of string
    """
    buf = infile.read(4)
    magic_num, = struct.unpack(b'<I', buf)
    if magic_num == 0xA1B2C3D4 or magic_num == 0x4D3C2B1A:
        return FileFormat.PCAP, buf
    elif magic_num == 0x0A0D0D0A:
        return FileFormat.PCAP_NG, buf
    else:
        return FileFormat.UNKNOWN, buf


def process_pcap_file(conn_dict, infile):
    file_format, head = get_file_format(infile)
    if file_format == FileFormat.PCAP:
        pcap_file = pcap.PcapFile(infile, head).read_packet
    elif file_format == FileFormat.PCAP_NG:
        pcap_file = pcapng.PcapngFile(infile, head).read_packet
    else:
        print("unknown file format.", file=sys.stderr)
        sys.exit(1)

    _filter = config.get_filter()
    for tcp_pac in packet_parser.read_package_r(pcap_file):
        # filter
        if not (_filter.by_ip(tcp_pac.source) or _filter.by_ip(tcp_pac.dest)):
            continue
        if not (_filter.by_port(tcp_pac.source_port) or _filter.by_port(tcp_pac.dest_port)):
            continue

        key = tcp_pac.gen_key()
        # we already have this conn
        if key in conn_dict:
            conn_dict[key].append(tcp_pac)
            # conn closed.
            if tcp_pac.pac_type == packet_parser.TcpPack.TYPE_CLOSE:
                conn_dict[key].finish()
                del conn_dict[key]

        # begin tcp connection.
        elif tcp_pac.pac_type == 1:
            conn_dict[key] = HttpConn(tcp_pac)
        elif tcp_pac.pac_type == 0:
            # tcp init before capture, we found a http request header, begin parse
            # if is a http request?
            if utils.is_request(tcp_pac.body):
                conn_dict[key] = HttpConn(tcp_pac)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("infile", nargs='?', help="the pcap file to parse")
    parser.add_argument("-i", "--ip", help="only parse packages with specified source OR dest ip")
    parser.add_argument("-p", "--port", type=int,
                        help="only parse packages with specified source OR dest port")
    parser.add_argument("-v", "--verbosity", help="increase output verbosity(-vv is recommended)",
                        action="count")
    parser.add_argument("-g", "--group", help="group http request/response by connection",
                        action="store_true")
    parser.add_argument("-o", "--output", help="output to file instead of stdout")
    parser.add_argument("-e", "--encoding", help="decode the data use specified encodings.")
    parser.add_argument("-b", "--beauty", help="output json in a pretty way.", action="store_true")
    parser.add_argument("-d", "--domain", help="filter http data by request domain")
    parser.add_argument("-u", "--uri", help="filter http data by request uri pattern")

    args = parser.parse_args()

    file_path = "-" if args.infile is None else args.infile

    _filter = config.get_filter()
    _filter.ip = args.ip
    _filter.port = args.port
    _filter.domain = args.domain
    _filter.uri_pattern = args.uri

    # deal with configs
    parse_config = config.get_config()
    if args.verbosity:
        parse_config.level = args.verbosity
    if args.encoding:
        parse_config.encoding = args.encoding
    parse_config.pretty = args.beauty
    parse_config.group = args.group

    if args.output:
        output_file = open(args.output, "w+")
    else:
        output_file = sys.stdout

    config.out = output_file

    conn_dict = OrderedDict()
    try:
        if file_path != '-':
            infile = io.open(file_path, "rb")
        else:
            infile = sys.stdin
        try:
            process_pcap_file(conn_dict, infile)
        finally:
            infile.close()
    finally:
        for conn in conn_dict.values():
            conn.finish()
        if args.output:
            output_file.close()


if __name__ == "__main__":
    try:
        main()
    except IOError as e:
        pass
