#!/usr/bin/env python
# Copyright (C) 2016 Maximilian Hils <git@maximilianhils.com>
# This file is part of HTTPReplay - http://jbremer.org/httpreplay/
# See the file 'LICENSE' for copying permission.

import click
import logging
import sys

from io import BytesIO

from httpreplay.reader import PcapReader
from httpreplay.smegma import TCPPacketStreamer
from httpreplay.cut import (
    http_handler, https_handler, smtp_handler
)
from httpreplay.misc import read_tlsmaster

try:
    from libmproxy import models
    from libmproxy.flow import FlowWriter
    from netlib.http import http1
except ImportError:
    try:
        from mitmproxy import models
        from mitmproxy.flow import FlowWriter
        from netlib.http import http1
    except ImportError:
        sys.exit(
            "In order to use this utility it is required to have the "
            "mitmproxy tool installed (`pip install httpreplay[mitmproxy]`)"
        )

logging.basicConfig(level=logging.INFO)

@click.command()
@click.argument("pcapfile", type=click.Path(file_okay=True))
@click.argument("mitmfile", type=click.File("wb"))
@click.option("--tlsmaster", type=click.Path(file_okay=True))
def convert(pcapfile, mitmfile, tlsmaster):
    if tlsmaster:
        tlsmaster = read_tlsmaster(tlsmaster)
    else:
        tlsmaster = {}

    handlers = {
        25: smtp_handler,
        80: http_handler,
        8000: http_handler,
        8080: http_handler,
        443: lambda: https_handler(tlsmaster),
        4443: lambda: https_handler(tlsmaster),
    }

    reader = PcapReader(pcapfile)
    reader.tcp = TCPPacketStreamer(reader, handlers)
    writer = FlowWriter(mitmfile)

    # Sort the http/https requests and responses by their timestamp.
    l = sorted(reader.process(), key=lambda x: x[1])
    for addrs, timestamp, protocol, sent, recv in l:
        if protocol not in ("http", "https"):
            continue

        srcip, srcport, dstip, dstport = addrs

        client_conn = models.ClientConnection.make_dummy((srcip, srcport))
        client_conn.timestamp_start = timestamp

        server_conn = models.ServerConnection.make_dummy((dstip, dstport))
        server_conn.timestamp_start = timestamp

        flow = models.HTTPFlow(client_conn, server_conn)

        # We need to manually read request and response bodies as the PCAP may
        # be incomplete and we'd complain on an unexpected body length.
        req_io = BytesIO(sent.raw)
        request = http1.read_request_head(req_io)
        request_body_size = http1.expected_http_body_size(request)

        if request_body_size > 0:
            request_body_size = -1

        request.data.content = \
            "".join(http1.read_body(req_io, request_body_size, None))

        flow.request = models.HTTPRequest.wrap(request)
        flow.request.host, flow.request.port = dstip, dstport
        flow.request.scheme = protocol

        resp_io = BytesIO(recv.raw)
        response = http1.read_response_head(resp_io)
        response_body_size = http1.expected_http_body_size(request, response)

        if response_body_size > 0:
            response_body_size = -1

        response.data.content = \
            "".join(http1.read_body(resp_io, response_body_size, None))

        flow.response = models.HTTPResponse.wrap(response)

        writer.add(flow)

if __name__ == "__main__":
    convert()
