#!/usr/bin/env python
import sys
import struct
import os
import threading
import json
import re
import urllib2
import uuid
from netlib.odict import ODictCaseless
from subprocess import call
from libmproxy import proxy, flow
from libmproxy.protocol.http import HTTPResponse
from libmproxy.proxy.config import ProxyConfig
from libmproxy.proxy.server import ProxyServer
from libmproxy.encoding import decode_gzip

from geventwebsocket import WebSocketServer, WebSocketApplication, Resource

sharedVars = {}
urlsToProxy = []
webSockets = []

def send_message(message):
    message = '{"msg":"' + message + '"}'
    sys.stdout.write(struct.pack('I', len(message)))
    sys.stdout.write(message)
    sys.stdout.flush()

class ChromeComm(WebSocketApplication):
    def on_open(self):
        send_message('Connection opened')
        webSockets.append(self.ws)
        self.ws.send(json.dumps({'method': 'rule-list', 'rules': urlsToProxy}))

    def on_message(self, message):
        if message == None:
            send_message('Disconnected')
        else:
            message = json.loads(message)
            if (message['method'] == 'add-rule'):
                ruleAlreadyAdded = False

                for url in urlsToProxy:
                    if (url['url'] == message['url']):
                        if (url['isLoaded'] == True):
                            fullFilePath = os.path.dirname(os.path.realpath(__file__)) + '/replacements/' + url['cachedFilename']
                            call(['open', fullFilePath])
                        ruleAlreadyAdded = True

                        self.ws.send(json.dumps({'method': 'rule-added', 'url': message['url']}))

                if not ruleAlreadyAdded:
                    send_message('Adding rule for: ' + message['url'])

                    # TODO guess file ext by Content-Type
                    regExpRes = re.search('.*\/(.*)', message['url'])
                    if (not regExpRes or regExpRes.group(1) == ''):
                        urlFilename = '.chromeproxy'
                    else:
                        urlFilename = '.' + regExpRes.group(1)
                        if (urlFilename.find('?') > -1):
                            urlFilename = urlFilename[:urlFilename.find('?')]
                        if urlFilename == '.':
                            urlFilename = '.chromeproxy'

                    while True:
                        filename = str(uuid.uuid1()) + urlFilename
                        fullFilePath = os.path.dirname(os.path.realpath(__file__)) + '/replacements/' + filename
                        if not os.path.isfile(fullFilePath):
                            break

                    localFile = open(fullFilePath, 'w')
                    localFile.write(message['responseContent'].encode('utf8'))
                    localFile.close()

                    call(['open', fullFilePath])
                    rule = {
                        'isEnabled': True,
                        'isLoaded': True,
                        'url': message['url'],
                        'responseHeaders': message['responseHeaders'],
                        'cachedFilename': filename
                    }

                    urlsToProxy.append(rule)

                    self.ws.send(json.dumps({'method': 'rule-added', 'rule': rule}))
                    save_request_map()


                    # send_message(message['responseContent'])
                    # urlsToProxy.append({
                    #     'isLoaded': False,
                    #     'url': message['url'],
                    #     'ws': self.ws
                    # });
                    # opener = urllib2.build_opener()
                    # send_message(map((lambda header: (header['name'], header['value'])), message['headers']))
                    # opener.addheaders.append(map((lambda header: (header['name'], header['value'])), message['headers']))
                    # f = opener.open(message['url'])
                    # urllib2.urlopen(message['url'])

            elif (message['method'] == 'remove-rule'):
                for url in urlsToProxy:
                    if (url['url'] == message['url']):
                        urlsToProxy.remove(url)
                        send_message('Removing rule for: ' + message['url'])
                        
                        self.ws.send(json.dumps({'method': 'rule-removed', 'rule': url}))
                        save_request_map();
                        break

            elif (message['method'] == 'udpate-rule'):
                for url in urlsToProxy:
                    if (url['url'] == message['url']):
                        url['isEnabled'] = message['isEnabled']
                        send_message('Updating rule for: ' + message['url'])

                        self.ws.send(json.dumps({'method': 'rule-updated', 'rule': url}))
                        save_request_map();
                        break

    def on_close(self, reason):
        send_message('bye')
        # TODO remove ws from webSockets list

def run_websocket():
    send_message('Running Websocket')
    WebSocketServer(
        ('', 8001),
        Resource({'/': ChromeComm})
    ).serve_forever()

class InjectingMaster(flow.FlowMaster):
    def __init__(self, server, state):
        flow.FlowMaster.__init__(self, server, state)

        relPath = os.path.dirname(os.path.realpath(__file__)) + '/'

        _requestMapFile = open(config['requestMapPath'], 'r+')
        requestMap = json.load(_requestMapFile)['rules']
        self._cachedFilesPath = relPath + 'replacements/'
        while len(urlsToProxy):
            urlsToProxy.pop()

        sharedVars['rules'] = [];

        for url in requestMap:
            urlsToProxy.append(url)
            sharedVars['rules'].append(url)


    def run(self):
        try:
            return flow.FlowMaster.run(self)
        except KeyboardInterrupt:
            self.shutdown()

    def handle_request(self, msg):
        f = flow.FlowMaster.handle_request(self, msg)

        send_message(f.request.scheme + '://' + f.request.host + f.request.path)
        for url in urlsToProxy:
            if (url['url'] == f.request.scheme + '://' + f.request.host + f.request.path and url['isEnabled'] == True):
                send_message('Serving cached file (' + url['cachedFilename'] + ')')
                localFile = open(self._cachedFilesPath + url['cachedFilename'], 'r');
                content = localFile.read()
                localFile.close();                

                # resp = HTTPResponse([1,1], 200, 'OK', ODictCaseless([['Content-Type', url['contentType']], ['via', 'chrome-proxy']]), content)
                responseHeaders = []
                hasViaHeader = False
                for header in url['responseHeaders']:
                    if (header['name'].lower() != 'content-encoding'):
                        if (header['name'].lower() == 'via'):
                            hasViaHeader = True
                            if (header['value'].find('chrome-proxy') == -1):
                                header['value'] += ', chrome-proxy'

                        responseHeaders.append((header['name'], header['value']))

                if (not hasViaHeader):
                    responseHeaders.append(['via', 'chrome-proxy'])

                # responseHeaders = filter((lambda header: header['name'] != 'Content-Encoding'), url['responseHeaders'])
                # responseHeaders = map((lambda header: (header['name'], header['value'])), responseHeaders)
                resp = HTTPResponse([1,1], 200, 'OK', ODictCaseless(responseHeaders), content)
                msg.reply(resp)
                break


        if 'Accept-Encoding' in f.request.headers:
            f.request.headers['Accept-Encoding'] = ['none']

        if f:
            msg.reply()

        return f

    def handle_responseheaders(self, f):
        # f = resp.flow

        isURLRequested = False
        for url in urlsToProxy:
            if (url['url'] == f.request.scheme + '://' + f.request.host + f.request.path and url['isLoaded'] == False):
                isURLRequested = True

        if not isURLRequested:
            f.response.stream = True

        f.reply()
        return f

    def handle_response(self, msg):
        f = flow.FlowMaster.handle_response(self, msg)

        # send_message(f.request.scheme + '://' + f.request.host + f.request.path)
        # for url in urlsToProxy:
        #     if (url['url'] == f.request.scheme + '://' + f.request.host + f.request.path and url['isLoaded'] == False):
        #         url['isLoaded'] = True
        #         while True:
        #             filename = str(uuid.uuid1()) + '.chromeproxy'
        #             fullFilePath = os.path.dirname(os.path.realpath(__file__)) + '/replacements/' + filename
        #             if not os.path.isfile(fullFilePath):
        #                 break

        #         localFile = open(fullFilePath, 'w')
        #         if (f.response.headers['Content-Encoding'] and f.response.headers['Content-Encoding'][0] == 'gzip'):
        #             localFile.write(decode_gzip(msg.content))
        #         else:
        #             localFile.write(msg.content)
        #         localFile.close()

        #         url['contentType'] = f.response.headers['Content-Type'][0]
        #         url['cachedFilename'] = filename

        #         call(['subl', '-n', fullFilePath])
        #         if url['ws']:
        #             send_message('sending message')
        #             url['ws'].send(json.dumps({'method': 'rule-added', 'url': url['url']}))

        #             del url['ws']

        #         save_request_map()


        if f:
            msg.reply()

        return f

config = {}

def save_request_map():
    requestMapFile = open(config['requestMapPath'], 'r+')
    requestMapFile.seek(0)
    requestMapFile.truncate()
    json.dump({'rules': urlsToProxy}, requestMapFile, sort_keys=True, indent=4)
    requestMapFile.close()

def main(argv):
    proxyConfig = proxy.config.ProxyConfig(
        port = 8889
    )

    relPath = os.path.dirname(os.path.realpath(__file__)) + '/'
    config['requestMapPath'] = relPath + 'requestmap.json'

    if not os.path.isdir(relPath + 'replacements/'):
        os.mkdir(relPath + 'replacements/')

    if not os.path.exists(config['requestMapPath']):
        f = open(config['requestMapPath'],'w')
        f.write('{"rules": []}')
        f.close()

    t1 = threading.Thread(target = run_websocket)
    t1.daemon = True
    t1.start()

    state = flow.State()
    server = proxy.server.ProxyServer(proxyConfig)

    m = InjectingMaster(server, state)
    m.run()

if __name__ == '__main__':
    main(sys.argv)
