Agent-Server Architecture

​

Overview

​

Components

​

Server

@asynccontextmanager
async def lifespan(app):
    # Initialise DB, session manager, and command queue on startup.
    global db, session_mgr, cmd_queue
    db          = Database()
    await db.__aenter__()
    session_mgr = SessionManager()
    cmd_queue   = CommandQueue()
    await session_mgr.restore_from_db(db)
    logger.info('server started', extra={'port': config.SERVER_PORT})

    yield  # server runs here

    # Shutdown — close DB connection cleanly
    if db:
        await db.__aexit__(None, None, None)
    logger.info('server stopped')

@app.post('/beacon')
async def beacon(request: Request) -> Response:
    # Handle all inbound agent messages: CHECKIN, TASK_PULL, TASK_RESULT
    source_ip = request.client.host if request.client else 'unknown'
    raw_body  = await request.body()

    # Reject oversized payloads early
    if len(raw_body) > MAX_BEACON_SIZE:
        logger.warning('payload too large', extra={
            'source_ip': source_ip,
            'payload_size': len(raw_body),
            'max_allowed': MAX_BEACON_SIZE,
        })
        return JSONResponse(status_code=413, content={'error': 'payload too large'})

• Accept and validate encrypted beacons from agents
• Manage agent sessions and track last-seen timestamps
• Queue and dispatch commands to agents
• Store task results and session data in SQLite database
• Perform nonce replay detection to prevent replay attacks

​

Agent

if __name__ == '__main__':
    try:
        check_lab_environment()
        BeaconLoop().run()
    except SystemExit:
        # check_lab_environment and TERMINATE signal both call sys.exit()
        raise
    except Exception as e:
        logger.error('catastrophic failure — agent exiting', extra={
            'reason':    str(e),
            'traceback': traceback.format_exc(),
        })
        sys.exit(1)

• Check in with the server and receive a unique session ID
• Send periodic beacons (TASK_PULL) to request new commands
• Execute received commands and capture stdout/stderr/exit code
• Report task results back to the server
• Implement exponential back-off on connection failures
• Terminate gracefully on TERMINATE signal

​

Request Flow

1. Initial Checkin
   • Agent sends CHECKIN message with system information (hostname, username, OS)
   • Server creates a new session, assigns a UUID, and stores it in the database
   • Server responds with the assigned session_id
2. Beacon Loop
   • Agent sleeps for configured interval (with jitter)
   • Agent sends TASK_PULL message to check for pending commands
   • Server checks the command queue for the session
   • Server responds with either:
     • TASK_DISPATCH containing a command to execute
     • TASK_PULL response with status: no_task
     • TERMINATE signal to shut down the agent
3. Task Execution
   • Agent executes the command using the executor module
   • Agent captures stdout, stderr, exit code, and duration
   • Agent sends TASK_RESULT message with execution data
   • Server stores the result and marks the task complete
4. Session Management
   • Server updates last_seen timestamp on every beacon
   • Operator can deactivate sessions, causing TERMINATE on next beacon
   • Sessions persist in the database across server restarts

​

Message Dispatch

async def _dispatch(msg_type: str, session_id: str,
                    payload: dict, source_ip: str) -> dict | None:
    # Route message to the correct handler and return the response payload dict

    if msg_type == mf.MSG_CHECKIN:
        return await _handle_checkin(payload, source_ip)

    if msg_type == mf.MSG_TASK_PULL:
        return await _handle_task_pull(session_id)

    if msg_type == mf.MSG_TASK_RESULT:
        return await _handle_task_result(session_id, payload)

    if msg_type == mf.MSG_HEARTBEAT:
        return await _handle_heartbeat(session_id)

    logger.warning('unknown msg_type', extra={'msg_type': msg_type})
    return None

​

Security Considerations

​

Scalability

• In-memory session state for fast lookups (restored from DB on startup)
• Async I/O with FastAPI and uvicorn for handling concurrent beacons
• SQLite database for persistence (can be replaced with PostgreSQL for production)
• Stateless beacon handling allows horizontal scaling behind a load balancer