import pdb

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

"""
# 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):
            try:
                data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
            except:
                print(obj)
                data = str(obj)[2:-1]
            return data

        return obj
"""

@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 flow dict
    def __init__(self, flow: dict):
        self.server_ip_address = flow["server_conn"]["ip_address"][0]
        self.tls = flow["server_conn"]["tls_established"]
        self.content = flow["request"]["content"]
        self.scheme = flow["request"]["scheme"]
        self.method = flow["request"]["method"]
        self.host = flow["request"]["host"]
        self.port = flow["request"]["port"]
        self.http_version = flow["request"]["http_version"]
        self.timestamp_start = flow["request"]["timestamp_start"]
        self.timestamp_end = flow["request"]["timestamp_end"]

        self.headers = []
        for k,v in flow["request"]["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 bFlowState(Enum):
    ERROR = 0
    UNSENT_HTTP_REQUEST = 1
    SENT_HTTP_REQUEST = 2
    UNSENT_HTTP_RESPONSE = 3
    SENT_HTTP_RESPONSE = 4

@dataclass
class bPacketType:
    NACK = 0
    ACK = 1
    KILL = 2
    WARNING = 3
    ERROR = 4
    PING = 5
    HTTP_REQUEST = 6
    HTTP_RESPONSE = 7

@dataclass
class bPacket:
    ptype: bPacketType
    flowid: str
    data: str

    def __init__(self, json: Dict):
        self.ptype = json["type"]
        self.flowid = str(json["id"])
        self.data = json["data"]

@dataclass
class FlowItem:
    state: bFlowState
    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: str, queue: Queue, path: str = None):
                threading.Thread.__init__(self)
                self.name = name
                # path
                self.path = path
                if not self.path:
                    self.path = os.environ.get("BIGSNITCH_PATH", None)
                if not self.path:
                    self.path = "tcp://127.0.0.1:12345"

                # queue for communicating with the main mitmproxy thread
                # contains tuples of (id, FlowItem)
                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
                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, flowitem = self.q.get(block=False)
                    if self.flows.get(i, None):
                        print(f"flow {i} doubled? ignoring...")
                        continue
                    else:
                        self.flows.append(flowitem)

                except Empty:
                    break

        def send_packet(self, pkg: bPacket):
            msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
            self.send(msg)

        def send_http_request(self, id: int, request: bRequest):
            pkg = bPacket(bPacketType.HTTP_REQUEST, id, request.json())
            self.send_packet(pkg)
            flows[id].state = bFlowState.SENT_HTTP_REQUEST
            flows[id].time = time.monotonic()

        def send_http_response(self, id: int, response: bResponse):
            pkg = bPacket(bPacketType.HTTP_RESPONSE, id, response.json())
            self.send_packet(pkg)
            flows[id].state = bFlowState.SENT_HTTP_RESPONSE
            flows[id].time = time.monotonic()

        # update all current flows
        # handles the state machine for each flow
        def update_flows(self):
            for id, flow in self.flows.items():
                # send the request if not sent
                if state == bFlowState.UNSENT_HTTP_REQUEST:
                    self.send_request(id, bRequest(flow.flow))

                elif state == bFlowState.SENT_HTTP_REQUEST:
                    # check timer, try resend
                    delta = time.monotonic() - flows[id].time
                    if delta > 5:
                        self.send_request(id, bRequest(flow.flow))
                        flows[id].retries -= 1
                        continue

                    if flows[id].retries <= 0:
                        flows[id].flow.kill()
                        print(f"http request {id} timed out! flow killed.")

                # send the response if not sent
                elif state == bFlowState.UNSENT_HTTP_RESPONSE:
                    self.send_response(id, bResponse(flow.flow))

                elif state == bFlowState.SENT_HTTP_RESPONSE:
                    # check timer, try resend
                    delta = time.monotonic() - flows[id].time
                    if delta > 5:
                        self.send_response(id, bResponse(flow.flow))
                        flows[id].retries -= 1
                        continue

                    if flows[id].retries <= 0:
                        flows[id].flow.kill()
                        print(f"http response {id} timedout! flow killed.")

                elif state == bFlowState.ERROR:
                    print(f"error in flow {id}!")

        # handle incoming packets / update the statemachine
        def handle_packets(self):
            while((self.socket.poll(50) & zmq.POLLIN) != 0):
                msg = self.socket.recv()
                try:
                    if msg:
                        result = json.loads(msg)
                        pkg = bPacket(json=result)
                        # flow ACKed
                        if pkg.ptype == bPacketType.ACK:
                            continue
                        # flow killed
                        elif pkg.ptype == bPacketType.KILL:
                            continue
                        else:
                            print(f"got unexpected message {pkg.ptype}")
                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):
                print("reconnecting")
                self.disconnect()
                time.sleep(1)
                self.connect()

        def connect(self):
                self.socket = self.context.socket(zmq.PAIR)
                self.socket.connect(self.path)
                #self.send_msg_and_ack({"msg": "ping"})
                print("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
        """