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

__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


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, output_file):
        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
        self.out = output_file

        # 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.pac_type == -1:
            # end of connection
            if self.status == HttpConn.STATUS_RUNNING:
                self.status = HttpConn.STATUS_CLOSED
            else:
                self.status = HttpConn.STATUS_ERROR

        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)

    def finish(self):
        self.http_parser.finish()
        result = self.processor.getvalue()
        print(result.encode('utf8'), file=self.out)
        self.out.flush()


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, ip, output_file, port):
    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)
    for tcp_pac in packet_parser.read_package_r(pcap_file):
        # filter
        if port is not None and tcp_pac.source_port != port and tcp_pac.dest_port != port:
            continue
        if ip is not None and tcp_pac.source != ip and tcp_pac.dest != ip:
            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, output_file)
        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, output_file)


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("-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")

    args = parser.parse_args()

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

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

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

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


if __name__ == "__main__":
    main()
