import grpc
import time
import os
import agent_pb2
import agent_pb2_grpc
import threading
import subprocess
import json
import platform
import jwt
import datetime
import hmac
import hashlib
SECRET_KEY = "cortex-secret-shared-key"
# --- Sandbox Policy Configuration ---
SANDBOX_POLICY = {
"MODE": "PERMISSIVE", # Toggle between "STRICT" and "PERMISSIVE"
"ALLOWED_COMMANDS": ["ls", "grep", "cat", "pwd", "git", "echo", "python", "whoami", "uname"],
"SENSITIVE_COMMANDS": ["rm", "cp", "mv", "chmod", "chown", "pkill"],
"DENIED_COMMANDS": ["sudo", "mkfs", "dd", "sh", "bash", "zsh"],
"WORKING_DIR": os.getcwd()
}
class AgentNode:
def verify_sandbox_policy(self, command_str):
"""Verifies if a command is allowed under the current sandbox policy."""
parts = command_str.strip().split()
if not parts:
return False, "Empty command"
base_cmd = parts[0]
# 1. ALWAYS Block Denied List (Strictly forbidden in all modes)
if base_cmd in SANDBOX_POLICY["DENIED_COMMANDS"]:
return False, f"Command '{base_cmd}' is strictly FORBIDDEN."
# 2. Path Guard (Simple string check for escaping ..)
if ".." in command_str:
return False, "Path traversal attempt detected (.. disallowed)."
# 3. Mode-specific Logic
if SANDBOX_POLICY["MODE"] == "STRICT":
# In STRICT mode, we check against the whitelist
if base_cmd not in SANDBOX_POLICY["ALLOWED_COMMANDS"] and base_cmd not in SANDBOX_POLICY["SENSITIVE_COMMANDS"]:
return False, f"STRICT MODE: Command '{base_cmd}' is NOT whitelisted."
# 4. Sensitive / Consent Check (Applied to both modes)
if base_cmd in SANDBOX_POLICY["SENSITIVE_COMMANDS"]:
return True, "SENSITIVE_CONSENT_REQUIRED"
return True, "OK"
def create_registration_token(self):
payload = {
"sub": "agent-node-007",
"workspace_id": "ws-production-001",
"iat": datetime.datetime.utcnow(),
"exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=10)
}
return jwt.encode(payload, SECRET_KEY, algorithm="HS256")
def __init__(self, node_id="agent-node-007"):
self.node_id = node_id
# Load certificates for mTLS
print("[🔐] Loading mTLS certificates...")
try:
with open('certs/client.key', 'rb') as f:
private_key = f.read()
with open('certs/client.crt', 'rb') as f:
certificate_chain = f.read()
with open('certs/ca.crt', 'rb') as f:
root_certificates = f.read()
# Create secure channel credentials
credentials = grpc.ssl_channel_credentials(
root_certificates=root_certificates,
private_key=private_key,
certificate_chain=certificate_chain
)
# Connect to localhost:50051 using secure channel
self.channel = grpc.secure_channel('localhost:50051', credentials)
self.stub = agent_pb2_grpc.AgentOrchestratorStub(self.channel)
print(f"[*] Agent Node {self.node_id} initialized with secure mTLS channel.")
except FileNotFoundError as e:
print(f"[!] Error: Certificates not found. Ensure generate_certs.sh was run. | {e}")
sys.exit(1)
# ... (rest of methods)
if __name__ == '__main__':
# We'll use a queue-based generator for better concurrency support
import queue
import sys
msg_queue = queue.Queue()
node = AgentNode()
# 1. Registration (Pre-handshake credentials with JWT)
token = node.create_registration_token()
reg = agent_pb2.NodeMessage(
registration=agent_pb2.RegistrationRequest(
node_id=node.node_id,
version="1.2.0",
platform=platform.system() + "-" + platform.machine(),
capabilities={"shell": True, "browser": False, "secure": True},
auth_token=token
)
)
msg_queue.put(reg)
def heartbeat_thread():
while True:
time.sleep(10)
hb = agent_pb2.NodeMessage(
heartbeat=agent_pb2.Heartbeat(
node_id=node.node_id,
cpu_usage_percent=1.2,
active_task_count=0
)
)
msg_queue.put(hb)
threading.Thread(target=heartbeat_thread, daemon=True).start()
def generator():
while True:
msg = msg_queue.get()
yield msg
try:
responses = node.stub.Connect(generator())
for response in responses:
payload_type = response.WhichOneof('payload')
if payload_type == 'registration_ack':
ack = response.registration_ack
if ack.success:
print(f"[*] Registered successfully. Session: {ack.session_id}")
else:
print(f"[!] Registration REJECTED: {ack.error_message}")
sys.exit(1)
elif payload_type == 'task_request':
task = response.task_request
print(f"[*] Task Received: {task.task_id}. Verifying signature...")
# Verify payload signature
expected_sig = hmac.new(SECRET_KEY.encode(), task.payload_json.encode(), hashlib.sha256).hexdigest()
if hmac.compare_digest(task.signature, expected_sig):
print(f" [OK] Signature verified. Checking sandbox policy...")
payload = json.loads(task.payload_json)
cmd = payload.get("command")
# --- Sandbox Enforcement ---
allowed, status_msg = node.verify_sandbox_policy(cmd)
if not allowed:
print(f" [⛔] Sandbox Violation: {status_msg}")
tr = agent_pb2.NodeMessage(
task_response=agent_pb2.TaskResponse(
task_id=task.task_id,
status=agent_pb2.TaskResponse.ERROR,
stderr=f"SANDBOX_VIOLATION: {status_msg}",
trace_id=task.trace_id
)
)
msg_queue.put(tr)
continue
if status_msg == "SENSITIVE_CONSENT_REQUIRED":
# In production: Wait for UI prompt. In POC: Log and proceed with a warning tag.
print(f" [⚠️] Sensitive Command Encountered: {cmd}. Automated approval assumed in POC.")
# -------------------------------
print(f" [OK] Execution starts: {cmd}")
res = subprocess.run(cmd, shell=True, capture_output=True, text=True)
# Send result back
tr = agent_pb2.NodeMessage(
task_response=agent_pb2.TaskResponse(
task_id=task.task_id,
status=agent_pb2.TaskResponse.SUCCESS,
stdout=res.stdout,
stderr=res.stderr,
duration_ms=0,
trace_id=task.trace_id
)
)
msg_queue.put(tr)
else:
print(f" [FAIL] Invalid signature for Task {task.task_id}! REJECTING.")
except grpc.RpcError as e:
print(f"[!] RPC Error: {e.code()} | {e.details()}")
if e.code() == grpc.StatusCode.UNAVAILABLE:
print(" Is the server running and reachable?")
elif e.code() == grpc.StatusCode.UNAUTHENTICATED:
print(" Authentication failed. Check certificates.")
