#!/usr/bin/env python3
#

import pdb

from mitmproxy import ctx
from mitmproxy.flow import Flow
import threading
from queue import Queue, Empty
import time
import zmq
import json
from enum import Enum
from dataclasses import dataclass
from typing import List, Dict

# this method is used to convert flow states (generated with get_state()) to json
def convert_to_strings(obj):
        if isinstance(obj, dict):
            return {convert_to_strings(key): convert_to_strings(value)
                    for key, value in obj.items()}
        elif isinstance(obj, list) or isinstance(obj, tuple):
            return [convert_to_strings(element) for element in obj]
        elif isinstance(obj, bytes):
            data = ""
            try:
                data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
            except:
                print(obj)
                data = str(obj)[2:-1]
            return data

        return obj
#
# bigsnitch communication types
#

@dataclass
class bHeader:
    key: str
    value: str

@dataclass
class bRequest:
    server_ip_address: str

    tls: str
    content: str
    scheme: str
    method: str
    host: str
    port: int
    http_version: str
    path: str
    timestamp_start: float
    timestamp_end: float
    # [("Header","Data")]
    headers: List[bHeader]

    error: str

    # init from dict
    def __init__(self, flow: dict):
        self.server_ip_address = flow["server_ip_address"]
        self.tls = flow["server_conn"]["tls_established"]
        self.content = flow["content"]
        self.scheme = flow["scheme"]
        self.method = flow["method"]
        self.host = flow["host"]
        self.port = flow["port"]
        self.http_version = flow["http_version"]
        self.timestamp_start = flow["timestamp_start"]
        self.timestamp_end = flow["timestamp_end"]

        for k,v in flow["headers"]:
            self.headers.append(bHeader(k,v))

@dataclass
class bResponse:
    status_code: int
    http_version: str
    reason: str
    content: str
    timestamp_start: float
    timestamp_end: float
    # [("Header","Data")]
    headers: List[bHeader]

    def __init__(self, flow: dict):
        self.status_code = flow["status_code"]
        self.http_version = flow["http_version"]
        self.reason = flow["reason"]
        self.content = flow["content"]
        self.timestamp_start = flow["timestamp_start"]
        self.timestamp_end = flow["timestamp_end"]

        for k,v in flow["headers"]:
            self.headers.append(bHeader(k,v))

@dataclass
class bFlow:
    uid: str;
    request: bRequest
    response: bResponse

    def __init__(self, flow: dict):
        self.uid = flow["id"]
        self.request =  bRequest(flow["request"])
        self.response =  bRequest(flow["response"])

#
# Networkthread state machine types
#

@dataclass
class FlowState(Enum):
    UNSENT_REQ = 0
    SENT_REQ = 1
    UNSENT_RES = 2
    SENT_RES = 3

# current flow state in Mitmproxy
@dataclass
class MitmState(Enum):
    ERROR = 0
    REQUESTHEADERS = 1
    REQUEST = 2
    RESPONSEHEADERS = 3
    RESPONSE = 4

# for use in NetworkThread queue
@dataclass
class FlowItem:
    id: int
    mitmstate: MitmState
    state: FlowState
    flow: Flow
    time: float = 0
    retries_left: int = 5

"""

The network thread communicates with the bigsnitch plugin using zeromq.

"""

class NetworkThread(threading.Thread):
        def __init__(self, name, queue):
                threading.Thread.__init__(self)
                self.name = name
                # queue for communicating with the main mitmproxy thread
                self.q = queue
                # for zmq use
                self.context = zmq.Context()
                # all current flows being handled by mitmproxy
                self.flows: Dict[FlowItem]
                # timer for sending pings to check if the connection broke
                self.timer = time.monotonic()
                # retries left for reconnecting / resending a broken flow
                self.retries = 5

        # send a single message, no checks involved
        def send(self, msg):
            a = convert_to_strings(msg)
            self.socket.send(str.encode(json.dumps(a)))

        # add new flows from the queue
        def get_new_flows(self):
            while True:
                try:
                    # get new flows that may occured
                    i, flow, typ = self.q.get(block=False)
                    if self.flows.get(i, None):
                        print(f"flow {i} doubled? ignoring...")
                        continue

                    # csave the new flows, if necessary
                    if typ == "request":
                        self.flows[i] = bFlow(FlowState.UNSENT_REQ, flow, time.monotonic(), 5)
                    elif typ == "response":
                        self.flows[i] = bFlow(FlowState.UNSENT_RES, flow, time.monotonic(), 5)
                except Empty:
                    break

        # update all current flows
        def update_flows(self):
            for k,v in self.flows.items():
                state, flow, timer, retries = v

                # state machine for flows

                if state == FlowState.UNSENT_REQ:
                    msg = b""
                    # send the request
                    self.send(msg)
                elif state == FlowState.SENT_REQ:
                    # check timer, try resend
                    pass
                elif state == FlowState.UNSENT_RES:
                    pass
                elif state == FlowState.SENT_RES:
                    pass

        def handle_packets(self):
            while((self.socket.poll(50) & zmq.POLLIN) != 0):
                msg = self.socket.recv()
                try:
                    if msg:
                        result = json.loads(msg)
                        # packet statemachine
                        if result["msg"] == "ack":
                            print("m ack received")
                            return result
                        else:
                            print("got unexpected message {result}")
                except json.JSONDecodeError:
                    print(f"malformed message received {msg}")

        def run(self):
                print("thread started")
                self.connect()
                while True:
                    self.timer = time.monotonic()
                    self.get_new_flows()
                    self.handle_packets()
                    self.update_flows()

                    if self.timer - time.monotonic() < -5:
                        pass
                        #self.send_msg_and_ack({"msg": "ping"})

        def disconnect(self):
                self.socket.setsockopt(zmq.LINGER,0)
                self.socket.close()
                print("disconnected")
                
        def reconnect(self):
                self.disconnect()
                time.sleep(1)
                self.connect()
        
        def connect(self):
                self.socket = self.context.socket(zmq.PAIR)
                self.socket.connect("tcp://127.0.0.1:12345")
                self.send_msg_and_ack({"msg": "ping"})
                print("successfully connected")

        def send_msg_and_ack(self, msg):
                self.timer = time.monotonic()
                while True:
                        #print("m sending")
                        self.send(msg)
                        if (self.socket.poll(50) & zmq.POLLIN) != 0:
                                msg = self.socket.recv()
                                try:
                                        if msg:
                                                result = json.loads(msg)
                                                if result["msg"] == "ack":
                                                        print("m ack received")
                                                        return result
                                                else:
                                                        print("got unexpected message {result}")
                                        
                                except json.JSONDecodeError:
                                        print(f"malformed message received {msg}")
                        print("no ack received, reconnecting...")
                        self.reconnect()
                return NO_MSG

class BigSnitchBridge:
    def __init__(self):
        self.q = Queue()
        self.thread = NetworkThread("network", self.q)
        self.thread.start()
    
    def request(self, flow):
        self.q.put_nowait((flow.id, flow, "request"))
        # intercept until ACK received
        flow.intercept()

    def response(self, flow):
        self.q.put_nowait((flow.id, flow, "response"))
        # intercept until ACK received
        flow.intercept()

    def error(self, flow):
        self.q.put_nowait((flow.id, flow, "error"))


"""
    def requestheaders(self, flow):
        self.q.put({'msg': 'requestheaders', 'flow': flow.get_state()})

    def responseheaders(self, flow):
        self.q.put({'msg': 'responseheaders', 'flow': flow.get_state()})

    def websocket_handshake(self):
        self.q.put({'msg': 'websocket_handshake', 'flow': flow.get_state()})

    def websocket_start(self, flow):
        self.q.put({'msg': 'websocket_start', 'flow': flow.get_state()})

    def websocket_message(self, flow):
        self.q.put({'msg': 'websocket_message', 'flow': flow.get_state()})

    def websocket_error(self, flow):
        self.q.put({'msg': 'websocket_error', 'flow': flow.get_state()})

    def websocket_end(self, flow):
        self.q.put({'msg': 'websocket_end', 'flow': flow.get_state()})
"""

addons = [
        BigSnitchBridge()
        ]