#!/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, textutils
from pcapparser.constant import FileFormat
from pcapparser.printer import HttpPrinter
from collections import OrderedDict
import struct

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


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, parse_config):
        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), parse_config)
        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 textutils.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"""
    buf = infile.read(4)
    infile.seek(0)
    magic_num, = struct.unpack(b'<I', buf)
    if magic_num == 0xA1B2C3D4 or magic_num == 0x4D3C2B1A:
        return FileFormat.PCAP
    elif magic_num == 0x0A0D0D0A:
        return FileFormat.PCAP_NG
    else:
        return FileFormat.UNKNOWN


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("infile", 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 = args.infile
    port = args.port
    ip = args.ip

    parse_config = ParseConfig()
    if args.verbosity:
        parse_config.level = args.verbosity
    if args.encoding:
        parse_config.encoding = args.encoding
    parse_config.pretty = args.beauty

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

    conn_dict = OrderedDict()
    try:
        with io.open(file_path, "rb") as infile:
            file_format = get_file_format(infile)
            if file_format == FileFormat.PCAP:
                pcap_file = pcap.PcapFile(infile).read_packet
            elif file_format == FileFormat.PCAP_NG:
                pcap_file = pcapng.PcapngFile(infile).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, parse_config)
                elif tcp_pac.pac_type == 0:
                    # tcp init before capture, we found a http request header, begin parse
                    # if is a http request?
                    if textutils.is_request(tcp_pac.body):
                        conn_dict[key] = HttpConn(tcp_pac, output_file, parse_config)
    finally:
        for conn in conn_dict.values():
            conn.finish()
        if args.output:
            output_file.close()


if __name__ == "__main__":
    main()
