#!/usr/bin/env python3
import pdb
from mitmproxy import ctx
import threading
from queue import Queue, Empty
import time
import zmq
import json
from enum import Enum

# 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 Request type
class bRequest:
    server_ip_address = ""

    tls = ""
    content = ""
    scheme = ""
    method = ""
    host = ""
    port = 0
    http_version = ""
    path = ""
    timestamp_start = 0.0
    timestamp_end = 0.0
    # [("Header","Data")]
    headers = []

    error = ""

class bResponse:
    status_code = 0
    http_version = ""
    reason = ""
    content = ""
    timestamp_start = 0.0
    timestamp_end = 0.0
    # [("Header","Data")]
    headers = []

class bFlow:
    uid = "";
    request = None
    response = None

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

"""

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
                self.q = queue
                # id : (state, flow, timer, retries left)
                self.flows = {}
                self.context = zmq.Context()
                # timer for sending pings
                self.timer = time.monotonic()
                # retries left for reconnecting
                self.retries = 5

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

                    if typ == "request":
                        flows[i] = (FlowState.UNSENT_REQ, flow, self.timer, 5)
                    elif typ == "response":
                        flows[i] = (FlowState.UNSENT_RES, flow, self.timer, 5)
                except Empty:
                    break

        # state machine for flows
        def handle_flow(self, state, flow):
            for i, v in flows.items:
                state, flow, timer, retries = v

                if state == FlowState.UNSENT_REQ:
                    # send it
                    a = convert_to_strings(msg)
                    self.socket.send(str.encode(json.dumps(a)))
                    pass
                elif state == FlowState.SENT_REQ:
                    # check timer, try resend
                    pass
                elif state == FlowState.UNSENT_RES:
                    pass
                elif state == FlowState.SENT_RES:
                    pass

                self.send_msg_and_ack(a)

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

                    if self.timer - time.monotonic() < -5:
                        self.send_msg_and_ack({"msg": "ping"})
                        
        def send(self, msg):
            a = convert_to_strings(msg)
            self.socket.send(str.encode(json.dumps(a)))

        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")
                        a = convert_to_strings(msg)
                        self.socket.send(str.encode(json.dumps(a)))
                        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()
        ]