import grpc
from concurrent import futures
import time
import agent_pb2
import agent_pb2_grpc
import queue
import threading
import jwt
import hmac
import hashlib
# In production, these would be in .env
SECRET_KEY = "cortex-secret-shared-key"
class AgentOrchestratorServicer(agent_pb2_grpc.AgentOrchestratorServicer):
def __init__(self):
self.nodes = {} # node_id -> queue for messages to node
def Connect(self, request_iterator, context):
node_id = None
# We need a way to send messages to the client from another thread/input
# In a real app, this would be triggered by an AI task planner
send_queue = queue.Queue()
def stream_messages():
while context.is_active():
try:
msg = send_queue.get(timeout=1.0)
yield msg
except queue.Empty:
continue
# Start a thread to handle incoming messages from the client
incoming_thread = threading.Thread(target=self._handle_incoming, args=(request_iterator, context, send_queue))
incoming_thread.start()
# Yield messages from the queue to the client
for msg in stream_messages():
yield msg
def _handle_incoming(self, request_iterator, context, send_queue):
try:
for message in request_iterator:
payload_type = message.WhichOneof('payload')
if payload_type == 'registration':
reg = message.registration
print(f"[*] Node {reg.node_id} registration request. Verifying token...")
try:
# Verify JWT
# In real app, we check if node_id matches token subject
decoded = jwt.decode(reg.auth_token, SECRET_KEY, algorithms=["HS256"])
print(f" [OK] Token verified for workspace: {decoded.get('workspace_id')}")
# Send ACK
ack = agent_pb2.ServerMessage(
registration_ack=agent_pb2.RegistrationResponse(
success=True,
session_id="session-secure-123"
)
)
send_queue.put(ack)
# Test 1: Allowed Command
t1_payload = '{"command": "whoami"}'
t1_sig = hmac.new(SECRET_KEY.encode(), t1_payload.encode(), hashlib.sha256).hexdigest()
send_queue.put(agent_pb2.ServerMessage(
task_request=agent_pb2.TaskRequest(
task_id="task-001-ALLOWED",
task_type="shell",
payload_json=t1_payload,
trace_id="trace-001",
signature=t1_sig
)
))
# Test 2: Sensitive Command (Consent Required)
t2_payload = '{"command": "rm -rf /tmp/node-test"}'
t2_sig = hmac.new(SECRET_KEY.encode(), t2_payload.encode(), hashlib.sha256).hexdigest()
send_queue.put(agent_pb2.ServerMessage(
task_request=agent_pb2.TaskRequest(
task_id="task-002-SENSITIVE",
task_type="shell",
payload_json=t2_payload,
trace_id="trace-002",
signature=t2_sig
)
))
# Test 3: Path Traversal Attempt (JAILBREAK)
t3_payload = '{"command": "cat ../.env.gitbucket"}'
t3_sig = hmac.new(SECRET_KEY.encode(), t3_payload.encode(), hashlib.sha256).hexdigest()
send_queue.put(agent_pb2.ServerMessage(
task_request=agent_pb2.TaskRequest(
task_id="task-003-TRAVERSAL",
task_type="shell",
payload_json=t3_payload,
trace_id="trace-003",
signature=t3_sig
)
))
# Test 4: STRICTLY Forbidden
t4_payload = '{"command": "sudo apt update"}'
t4_sig = hmac.new(SECRET_KEY.encode(), t4_payload.encode(), hashlib.sha256).hexdigest()
send_queue.put(agent_pb2.ServerMessage(
task_request=agent_pb2.TaskRequest(
task_id="task-004-FORBIDDEN",
task_type="shell",
payload_json=t4_payload,
trace_id="trace-004",
signature=t4_sig
)
))
# Test 5: Non-whitelisted but Allowed in PERMISSIVE
t5_payload = '{"command": "df -h"}'
t5_sig = hmac.new(SECRET_KEY.encode(), t5_payload.encode(), hashlib.sha256).hexdigest()
send_queue.put(agent_pb2.ServerMessage(
task_request=agent_pb2.TaskRequest(
task_id="task-005-PERMISSIVE",
task_type="shell",
payload_json=t5_payload,
trace_id="trace-005",
signature=t5_sig
)
))
print("[*] Sequence of 5 test tasks dispatched to verify Sandbox Policy (PERMISSIVE mode).")
except jwt.ExpiredSignatureError:
print(f" [FAIL] Token for {reg.node_id} expired.")
ack = agent_pb2.ServerMessage(
registration_ack=agent_pb2.RegistrationResponse(
success=False,
error_message="Authentication token expired."
)
)
send_queue.put(ack)
except jwt.InvalidTokenError as e:
print(f" [FAIL] Invalid token for {reg.node_id}: {e}")
ack = agent_pb2.ServerMessage(
registration_ack=agent_pb2.RegistrationResponse(
success=False,
error_message=f"Invalid authentication token: {e}"
)
)
send_queue.put(ack)
elif payload_type == 'heartbeat':
hb = message.heartbeat
# print(f"[+] Heartbeat from {hb.node_id}: CPU {hb.cpu_usage_percent}%")
pass
elif payload_type == 'task_response':
res = message.task_response
print(f"[!] Task Finished: {res.task_id} | Status: {res.status}")
print(f" Stdout: {res.stdout.strip()}")
except Exception as e:
print(f"[!] Error handling incoming stream: {e}")
def serve():
# Load certificates for mTLS
print("[🔐] Loading mTLS certificates...")
with open('certs/server.key', 'rb') as f:
private_key = f.read()
with open('certs/server.crt', 'rb') as f:
certificate_chain = f.read()
with open('certs/ca.crt', 'rb') as f:
root_certificates = f.read()
# Create server credentials
# require_client_auth=True enforces bidirectional verification (mTLS)
server_credentials = grpc.ssl_server_credentials(
[(private_key, certificate_chain)],
root_certificates=root_certificates,
require_client_auth=True
)
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
agent_pb2_grpc.add_AgentOrchestratorServicer_to_server(AgentOrchestratorServicer(), server)
# Use secure_port instead of insecure_port
server.add_secure_port('[::]:50051', server_credentials)
print("[*] Cortex Secure Server POC listening on port 50051 (mTLS enabled)...")
server.start()
server.wait_for_termination()
if __name__ == '__main__':
serve()
