import grpc import time import os import agent_pb2 import agent_pb2_grpc import threading import subprocess import json import jwt import datetime import hmac import hashlib import queue import sys import platform from concurrent import futures SECRET_KEY = "cortex-secret-shared-key" class BaseSkill: """Interface for pluggable node capabilities.""" def execute(self, task, sandbox, on_complete): raise NotImplementedError def cancel(self, task_id): return False class ShellSkill(BaseSkill): """Default Skill: Executing shell commands.""" def __init__(self): self.processes = {} # task_id -> Popen self.lock = threading.Lock() def execute(self, task, sandbox, on_complete): try: payload = json.loads(task.payload_json) cmd = payload.get("command") allowed, status_msg = sandbox.verify(cmd) if not allowed: return on_complete(task.task_id, {"stderr": f"SANDBOX_VIOLATION: {status_msg}", "status": 2}, task.trace_id) print(f" [🐚] Executing Shell: {cmd}") p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) with self.lock: self.processes[task.task_id] = p timeout = task.timeout_ms / 1000.0 if task.timeout_ms > 0 else None stdout, stderr = p.communicate(timeout=timeout) on_complete(task.task_id, {"stdout": stdout, "stderr": stderr, "status": 1 if p.returncode == 0 else 2}, task.trace_id) except subprocess.TimeoutExpired: self.cancel(task.task_id) on_complete(task.task_id, {"stderr": "TIMEOUT", "status": 2}, task.trace_id) except Exception as e: on_complete(task.task_id, {"stderr": str(e), "status": 2}, task.trace_id) finally: with self.lock: self.processes.pop(task.task_id, None) def cancel(self, task_id): with self.lock: p = self.processes.get(task_id) if p: print(f"[🛑] Killing Shell Process: {task_id}") p.kill() return True return False class SkillManager: """Orchestrates multiple skills and manages the worker thread pool.""" def __init__(self, max_workers=5): self.executor = futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="skill-worker") self.active_tasks = {} # task_id -> future self.skills = {"shell": ShellSkill()} self.max_workers = max_workers self.lock = threading.Lock() def submit(self, task, sandbox, on_complete): with self.lock: if len(self.active_tasks) >= self.max_workers: return False, "Node Capacity Reached" # Determine Skill (Default to shell for now) skill = self.skills.get("shell") future = self.executor.submit(skill.execute, task, sandbox, on_complete) self.active_tasks[task.task_id] = future # Cleanup hook future.add_done_callback(lambda f: self._cleanup(task.task_id)) return True, "Accepted" def cancel(self, task_id): with self.lock: # Tell all skills to try and cancel this ID cancelled = any(s.cancel(task_id) for s in self.skills.values()) return cancelled def get_active_ids(self): with self.lock: return list(self.active_tasks.keys()) def _cleanup(self, task_id): with self.lock: self.active_tasks.pop(task_id, None) class SandboxEngine: def __init__(self): self.policy = None def sync(self, p): self.policy = {"MODE": "STRICT" if p.mode == agent_pb2.SandboxPolicy.STRICT else "PERMISSIVE", "ALLOWED": list(p.allowed_commands), "DENIED": list(p.denied_commands), "SENSITIVE": list(p.sensitive_commands)} def verify(self, command_str): if not self.policy: return False, "No Policy" parts = (command_str or "").strip().split() if not parts: return False, "Empty" base_cmd = parts[0] if base_cmd in self.policy["DENIED"]: return False, "Forbidden" if self.policy["MODE"] == "STRICT" and base_cmd not in self.policy["ALLOWED"]: return False, "Not Whitelisted" return True, "OK" class AgentNode: def __init__(self, node_id="agent-node-007"): self.node_id = node_id self.skills = SkillManager() self.sandbox = SandboxEngine() self.task_queue = queue.Queue() # gRPC Setup with open('certs/client.key', 'rb') as f: pkey = f.read() with open('certs/client.crt', 'rb') as f: cert = f.read() with open('certs/ca.crt', 'rb') as f: ca = f.read() creds = grpc.ssl_channel_credentials(ca, pkey, cert) self.channel = grpc.secure_channel('localhost:50051', creds) self.stub = agent_pb2_grpc.AgentOrchestratorStub(self.channel) def _create_token(self): return jwt.encode({"sub": self.node_id, "iat": datetime.datetime.utcnow(), "exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=10)}, SECRET_KEY, algorithm="HS256") def sync_configuration(self): print(f"[*] Handshake: {self.node_id}") reg = agent_pb2.RegistrationRequest(node_id=self.node_id, auth_token=self._create_token(), node_description="Refactored Stateful Node", capabilities={"shell": "v1"}) res = self.stub.SyncConfiguration(reg) if res.success: self.sandbox.sync(res.policy); print("[OK] Policy Synced.") else: print(f"[!] Rejected: {res.error_message}"); sys.exit(1) def start_health_reporting(self): def _gen(): while True: ids = self.skills.get_active_ids() yield agent_pb2.Heartbeat(node_id=self.node_id, cpu_usage_percent=1.0, active_worker_count=len(ids), max_worker_capacity=5, running_task_ids=ids) time.sleep(10) threading.Thread(target=lambda: list(self.stub.ReportHealth(_gen())), daemon=True).start() def run_task_stream(self): def _gen(): while True: yield self.task_queue.get() responses = self.stub.TaskStream(_gen()) for msg in responses: kind = msg.WhichOneof('payload') if kind == 'task_request': self._handle_task(msg.task_request) elif kind == 'task_cancel': if self.skills.cancel(msg.task_cancel.task_id): self.task_queue.put(agent_pb2.ClientTaskMessage(task_response=agent_pb2.TaskResponse( task_id=msg.task_cancel.task_id, status=agent_pb2.TaskResponse.CANCELLED))) elif kind == 'work_pool_update': for tid in msg.work_pool_update.available_task_ids: if len(self.skills.get_active_ids()) < self.skills.max_workers: self.task_queue.put(agent_pb2.ClientTaskMessage(task_claim=agent_pb2.TaskClaimRequest(task_id=tid, node_id=self.node_id))) def _handle_task(self, task): # Sig Verify sig = hmac.new(SECRET_KEY.encode(), task.payload_json.encode(), hashlib.sha256).hexdigest() if not hmac.compare_digest(task.signature, sig): return print("[!] Sig Fail") self.skills.submit(task, self.sandbox, self._on_finish) def _on_finish(self, tid, res, trace): print(f"[*] Task {tid} finished.") status = agent_pb2.TaskResponse.SUCCESS if res['status'] == 1 else agent_pb2.TaskResponse.ERROR tr = agent_pb2.TaskResponse(task_id=tid, status=status, stdout=res.get('stdout',''), stderr=res.get('stderr',''), trace_id=trace) self.task_queue.put(agent_pb2.ClientTaskMessage(task_response=tr)) if __name__ == '__main__': node = AgentNode() node.sync_configuration() node.start_health_reporting() node.run_task_stream()