Build a Basic AI Agent From Scratch: Security III A new policy gate for AI agents adds three security layers—path scoping, a shell denylist, and an SSRF guard—to block dangerous tool calls before they reach execution, according to the latest installment of the 'Build a Basic AI Agent From Scratch' series. The gate runs before the mode-based permission check and applies to every tool that uses paths, commands, or URLs, enforcing resource limits, secret scrubbing, audit logging, and a kill switch. Previous parts of Build a Basic AI Agent From Scratch : You can find and clone this code in this blog series' Github repo . In the previous part we started closing the gaps left open by human-in-the-loop: a Docker sandbox to contain runaway commands, prompt-injection defenses so the model stops trusting tool output as instructions, and schema validation so malformed tool calls never reach execution. In this part we finish the job. We will harden the tool policy gate with path scoping, a shell denylist, and an SSRF guard, enforce resource and cost limits so a stuck loop cannot run forever, scrub secrets out of the container environment, and add audit logging and a kill switch so every decision is recorded and any session can be aborted mid-flight. In the Human in the Loop & Security part, check permission was a single function: read tools and planning tools were always allowed, write tools were allowed inside the working directory in acceptEdits , everything else asked. That was a mode-based decision. We now add a policy gate that runs before the mode decision. The gate has three layers: .. traversal and rejected if it escapes the project working directory. This stops the agent from touching ~/.ssh/id rsa , /etc/passwd , or anything outside the project tree. run bash command is screened against a regex denylist of dangerous patterns fork bombs, dd , mkfs , redirects into /etc/ , ... and a token-level denylist of binaries the agent should never invoke sudo , nc , curl , chmod , docker , ... . This catches destructive and exfiltration commands. webfetch hit internal services or the cloud metadata endpoint. The guard resolves the URL's host, checks the resulting IP against loopback, link-local, multicast, and RFC1918 private ranges, and blocks them. python def check tool policy tool name: str, args: dict str, Any , working dir: Path, - tuple bool, str | None : """Run all policy layers. Returns allowed, reason . Called by agent.check permission BEFORE the mode-based decision. A False here is a hard block that no mode can override. """ Layer 1: path scope. ok, reason = check path scope tool name, args, working dir if not ok: return False, reason Layer 2: shell policy. if tool name == "run bash": ok, reason = check shell policy args.get "command", "" if not ok: return False, reason Layer 3: web / SSRF policy. if tool name == "webfetch": ok, reason = check web policy args.get "url", "" if not ok: return False, reason return True, None The old check only ran on write tools. Now, the check is generalized to every tool that uses paths: read file , glob files , grep , write file , edit file . Each tool's path argument is resolved handling relative paths, symlinks, and .. traversal and rejected if it escapes the working directory: PATH TOOLS: dict str, str = { "read file": "path", "glob files": "path", "grep": "path", "write file": "path", "edit file": "path", } def check path scope tool name: str, args: dict str, Any , working dir: Path, - tuple bool, str | None : """Reject path-bearing tool calls whose target escapes working dir. Note: the Docker mount already constrains the container's view of the filesystem; this check is a defense-in-depth layer on the host side so a malicious path is rejected before it ever reaches docker. """ if tool name not in PATH TOOLS: return True, None raw = args.get PATH TOOLS tool name if not raw: return True, None missing arg is a schema problem, not a scope problem try: target = Path raw if not target.is absolute : target = working dir / target target.resolve .relative to working dir.resolve return True, None except ValueError, OSError, RuntimeError as e: return False, f"Path '{raw}' is outside the working directory " f" {working dir} . File tools may only touch paths inside " f"the project root. {e} " This is defense-in-depth on the host side. The Docker mount already constrained the container's view of the file system, but a malicious path is rejected here before it even reaches the sandbox. By far, run bash is the most dangerous tool, so it gets its own screening. Two checks run against every command. First, a regex denylist of dangerous patterns: SHELL DENYLIST PATTERNS = re.compile r"\brm\s+-rf?\s+ /|~|\ |\$HOME|\.\. ", re.I , "recursive delete of a broad or root target" , re.compile r" \s /etc/", re.I , "redirect into /etc/ system files " , re.compile r"\bmkfs\b", re.I , "filesystem format command" , re.compile r"\bdd\b\s+if=", re.I , "raw disk write via dd" , re.compile r":\ \ \s \{", re.I , "fork-bomb pattern" , re.compile r"\b eval|exec \b", re.I , "eval/exec in a shell command injection risk " , re.compile r" \s /dev/sd", re.I , "write to a block device" , re.compile r"\bhistory\s+-c\b", re.I , "history wipe" , re.compile r"\bexport\s+PATH=", re.I , "PATH override could shadow binaries " , Then, the command is lexically parsed and every token is checked against a deny list. The agent should never use any of those binaries in any command: SHELL DENYLIST BINARIES = frozenset { "docker", sandbox escape / host control "sudo", "su", privilege escalation "nc", "netcat", "ncat", reverse shells / exfil "curl", "wget", exfil / SSRF "chmod", "chown", permission tampering "mkfs", "dd", destructive disk ops "shutdown", "reboot", "halt", "poweroff", "systemctl", "service", "crontab", "at", } webfetch is screened against server-side request forgery. Cloud metadata endpoints, localhost, and private IP ranges are all blocked: WEB DENYLIST HOSTS = frozenset { "169.254.169.254", AWS / GCP / Azure cloud metadata "metadata.google.internal", GCP metadata "metadata.azure.com", Azure metadata "0.0.0.0", "::1", "localhost", } def check web policy url: str - tuple bool, str | None : """Reject URLs that target loopback / link-local / private ranges.""" from urllib.parse import urlparse try: parsed = urlparse url except ValueError as e: return False, f"Unparseable URL: {e}" if parsed.scheme not in "http", "https" : return False, f"Unsupported scheme '{parsed.scheme}'." host = parsed.hostname if not host: return False, "URL has no host component." if host.lower in WEB DENYLIST HOSTS: return False, f"Blocked host '{host}' loopback / metadata ." Resolve and check the IP family. try: infos = socket.getaddrinfo host, None except socket.gaierror: Let the actual fetcher surface the DNS error. return True, None for info in infos: ip = info 4 0 try: addr = ipaddress.ip address ip except ValueError: continue if addr.is loopback or addr.is link local or addr.is multicast: return False, f"Blocked IP '{ip}' loopback / link-local / multicast ." if addr.is private: return False, f"Blocked IP '{ip}' RFC1918 private range SSRF guard ." return True, None localhost and 127.0.0.1 are where services that aren't exposed to the public network live: databases Postgres on 5432, Redis on 6379 , admin panels, metrics endpoints, internal APIs, and sometimes cloud-metadata proxies. A prompt-injected webfetch "http://localhost:8080/admin/delete-all" could trigger destructive actions on a service that trusts local connections and therefore has no auth. Blocking loopback prevents the agent from being used as a pivot into the host's internal network. The host is resolved via getaddrinfo and each returned IP is checked with ipaddress . Loopback, link-local, multicast, and RFC1918 private ranges are all blocked. This stops the model from fetching http://169.254.169.254/latest/meta-data/iam/security-credentials/ to steal cloud credentials, or probing internal services on 10.0.0.0/8 . On top of the hard policy gate, some patterns are dangerous enough that they require explicit confirmation regardless of permission mode . These are the irreversible / exfiltration class: NOTE: for run bash, anything already covered by SHELL DENYLIST BINARIES or SHELL DENYLIST PATTERNS rm -rf, sudo/su, docker, chmod, curl/wget/nc/netcat/ncat is a hard block in layer 1 and never reaches this layer, so it is deliberately NOT duplicated here. Only patterns that layer 1 does not already hard-block belong in this list. ALWAYS CONFIRM ARG PATTERNS: list tuple str, re.Pattern str , str = force-push to git git itself is not on the shell denylist "run bash", re.compile r"\bgit\s+push\s+ -f|--force ", re.I , "force-push to git rewrites remote history " , overwriting a file with empty content = delete "write file", re.compile r'"content"\s :\s "\s "' , "write file with empty content on an existing file delete-via-empty " , The check permission function now has a 3-layer structure: check tool policy . Checks path scope, shell policy, SSRF guard. A False here blocks the call regardless of mode. always confirm required returns True, reason : dangerouslySkipPermissions : the call is default / acceptEdits : the user is prompted with a DESTRUCTIVE label. default / acceptEdits / dangerouslySkipPermissions logic, with the addition that write file emptying an existing file forces a DELETE-via-empty confirmation even in acceptEdits . python def check permission tool name: str, args: dict, mode: PermissionMode, working dir: Path, - tuple bool, str | None : Layer 1: hard policy gate path scope, shell, SSRF . ok, reason = tool policy.check tool policy tool name, args, working dir if not ok: return False, reason Layer 2: always-confirm & irreversible patterns. confirm required, confirm reason = tool policy.always confirm required tool name, args if confirm required: if mode == PermissionMode.DANGEROUSLY SKIP PERMISSIONS: return False, f"Blocked even in dangerouslySkipPermissions: " f"{confirm reason or 'irreversible action'}. " "Run this command manually outside the agent if it is truly intended." return ask permission f"{tool name} DESTRUCTIVE ", args , None Layer 3: mode-based decision read/planning free, acceptEdits for in-tree writes, etc. ... always confirm required returns bool, reason rather than a bare bool, so this layer carries its own specific reason instead of falling back on whatever layer 1 happened to return which, on the success path, is always None . check permission itself also returns allowed, reason , so rejections carry a machine-readable reason surfaced to the LLM and the audit log. One more control fits here: a --tools CLI flag that accepts a comma-separated allow list of tool names. The harness filters both the in-process registry and the schemas exposed to the LLM, so the model never even sees tools it can't call: bash $ python agent.py --tools read file,glob files,grep,run bash If we know our agent will never need to write files, we can reduce the danger blast radius enormously by limiting the allowed tools. We will add to the harness some resource and cost controls to limit the agent from going out of hand while we are not looking at it. IterationCaps tracks two counters: turns --max-turns . tool calls --max-tool-calls . @dataclass class IterationCaps: max turns per user msg: int = DEFAULT MAX TURNS PER USER MSG max tool calls per session: int = DEFAULT MAX TOOL CALLS PER SESSION turns: int = 0 tool calls: int = 0 def reset turn self - None: self.turns = 0 def bump turn self - str | None: self.turns += 1 if self.turns self.max turns per user msg: return f"Reached the per-turn iteration cap " f" {self.max turns per user msg} LLM turns . Stop calling " f"tools and give the user a concise summary of progress." return None When a cap is reached, the loop injects a "stop and summarize" message into messages and logs an iteration cap hit audit event with scope of per turn or session . The model is told to stop calling tools and give the user a concise summary. The defaults are chosen to be generous enough for real coding tasks but short enough that a stuck loop is killed quickly. ContextBudget tracks cumulative tokens we do a rough estimation of 4 chars per token and trims the message history before each LLM call: The 4-chars-per-token estimate is a rough heuristic, but it's good enough for a safety budget. Instead, we could get the actual token count from the specific model API, but to avoid implementing a method to get the actual token count for each different LLM Provider, we will use the rough estimation for now. @dataclass class ContextBudget: max tokens: int = DEFAULT MAX CONTEXT TOKENS trim threshold: float = 0.8 keep recent: int = DEFAULT KEEP RECENT MESSAGES def check and trim self, messages: list - tuple bool, str | None : """Trim messages in place if over budget.""" self.last estimate = self.estimate total messages if self.last estimate <= int self.max tokens self.trim threshold : return False, None We always keep messages 0 system prompt and the last keep recent messages. Everything in between is a candidate for trimming. if len messages <= self.keep recent + 1: return False, None cut start = 1 cut end = len messages - self.keep recent dropped = messages cut start:cut end summary = self. summarize dropped messages cut start:cut end = { "role": "system", "content": summary, } ... If the estimate exceeds max tokens trim threshold by default it's 80% , it replaces everything between the system prompt and the last n keep recent messages by default 8 messages with a summary message. The summary is deterministic and LLM-free: it records dropped message count, total chars, tool-call names, a SHA-256 hash of the dropped conversation first 16 hex , and an instruction to re-read files rather than relying on dropped context: php @staticmethod def summarize dropped: list - str: tool calls: list str = total chars = 0 for m in dropped: ... collect tool-call names and char counts digest = hashlib.sha256 blob.encode "utf-8", "replace" .hexdigest :16 lines = " context-trim summary ", f"Earlier conversation dropped to stay within the token budget.", f"Dropped messages: {len dropped } {total chars} chars .", f"Tools called in dropped section: {', '.join tool calls or 'none'}.", f"Conversation hash first 16 hex : {digest}.", "Re-read any files you need rather than relying on the dropped context.", return "\n".join lines A companion function, cap tool result , caps each tool result at 32 KB before insertion into messages , with a notice telling the model to use read file offset/limit for more. Without this, a single read file of a 50k-line file would blow the budget in one call. By default, the budget is 24,000 tokens, but it can be overwritten with --max-context-tokens . We already saw the per-tool-call timeout in the previous part. The LLM call itself now carries timeout=llm timeout default 120s, via --llm-timeout . A timeout or connection failure is caught and reported to the user rather than crashing the process, with an llm call error audit event: try: response = client.chat.completions.create model="gemma4", messages=messages, tools=active schemas, temperature=0.7, timeout=llm timeout, except Exception as e: audit.log "llm call error", error=str e print f"\n llm error {e}" print "Assistant: I could not reach the model. " "Please check the endpoint and retry." break CostTracker accumulates API spend. After each chat.completions.create , record usage reads response.usage or None for local backends like Ollama that don't report usage and accumulates total tokens in , total tokens out , and total cost usd : @dataclass class CostTracker: max cost usd: float = DEFAULT MAX COST USD price in: float = DEFAULT PRICE PER 1K IN price out: float = DEFAULT PRICE PER 1K OUT total tokens in: int = 0 total tokens out: int = 0 total cost usd: float = 0.0 calls: int = 0 def record usage self, usage: Any | None - None: self.calls += 1 if usage is None: return pt = getattr usage, "prompt tokens", None ct = getattr usage, "completion tokens", None ... handle dict or object form self.total tokens in += pt or 0 self.total tokens out += ct or 0 self.total cost usd = self.total tokens in / 1000.0 self.price in + self.total tokens out / 1000.0 self.price out def check self - str | None: if self.total cost usd = self.max cost usd: return f"Cost limit reached: ${self.total cost usd:.4f} = " f"${self.max cost usd:.4f} cap. Stop and report to the user." return None check returns a reason when spend ≥ max cost usd ; the loop logs a cost limit hit audit event and stops with a user-facing message. CLI flag: --max-cost-usd default $5 . The session-end summary now also prints the full cost breakdown calls , tokens in , tokens out , cost . The model can leak anything it can see. If secrets get in the context window of our agent they can easily leak via logs, or tool uses. So we will have to make sure that the model never sees secrets . To enforce this, we will use three layers: Two complementary checks run at startup. scan environment for secrets lists all host env vars matching the pattern KEY|SECRET|TOKEN|PASSWORD|PASSWD|CREDENTIAL|APIKEY|API KEY|AUTH case-insensitive . The count is logged in the audit config event as secret env count . SECRET ENV PATTERN = re.compile r" . ?:KEY|SECRET|TOKEN|PASSWORD|PASSWD|CREDENTIAL|APIKEY|API KEY|AUTH . ", re.IGNORECASE, def scan environment for secrets - list tuple str, str : """Return a list of name, source for env vars that look like secrets.""" found: list tuple str, str = for name in sorted os.environ : if SECRET ENV PATTERN.match name : found.append name, "env" return found audit system prompt prompt statically checks the system prompt template before it is sent to the model. It flags os.environ / os.getenv interpolation patterns in the template e.g. f"... {os.environ 'API KEY' } ..." and literal occurrences of host secret env-var names in the prompt. PROMPT INTERPOLATION PATTERN = re.compile r"\{. ?:os\.environ|os\.getenv|getenv|environ . \}", re.IGNORECASE, def audit system prompt prompt: str - list str : """Statically check prompt for patterns that could leak secrets.""" warnings: list str = if PROMPT INTERPOLATION PATTERN.search prompt : warnings.append "System prompt appears to interpolate os.environ / os.getenv. " "Never put secrets in the system prompt, the model can leak " "them in tool calls or responses." for name in os.environ: if SECRET ENV PATTERN.match name and name in prompt: warnings.append f"System prompt literally contains the env-var name " f"'{name}'. Even if this is the name and not the value, " f"its presence may prompt the model to look it up." return warnings If any warning is found, the harness refuses to start RuntimeError . This is a defense-in-depth check: the current prompt is a static string with no env-var interpolation, but if someone later edits it to inject a key, this catches it at startup before the model ever sees the prompt. The sandbox container only inherits an allow list of env vars from the host. Everything else is stripped out before the container starts: ALLOWED CONTAINER ENV = frozenset { "PATH", "HOME", "USER", "LANG", "LC ALL", "TERM", "AGENT SESSION ID", "AGENT SESSION TOKEN" } def build container env session id: str, session token: str - dict str, str : """Return the minimal environment dict for the sandbox container. Only ALLOWED CONTAINER ENV variables are inherited from the host; everything else including any secret-looking env vars is stripped. """ env: dict str, str = {} for name in ALLOWED CONTAINER ENV: if name in os.environ: env name = os.environ name env "AGENT SESSION ID" = session id env "AGENT SESSION TOKEN" = session token return env DockerSandbox. init accepts a container env dict; start container passes each entry as a -e NAME=VALUE flag to docker run . A companion check lists host paths that must never be bind-mounted into the container: ~/.aws , ~/.ssh , ~/.config/gcloud , ~/.docker , ~/.netrc , ~/.kube , ~/.gnupg . check credential mounts reports which exist on the host so the operator can verify the docker run command never mounts them: CREDENTIAL MOUNT PATHS: list Path = Path.home / ".aws", Path.home / ".ssh", Path.home / ".config" / "gcloud", Path.home / ".docker", Path.home / ".netrc", Path.home / ".kube", Path.home / ".gnupg", def check credential mounts - list Path : """Return the subset of credential paths that exist on the host.""" return p for p in CREDENTIAL MOUNT PATHS if p.exists SessionCredentials generates fresh per-session credentials: php class SessionCredentials: def init self - None: self.session id = secrets.token urlsafe 12 self.session token = secrets.token urlsafe 32 self. revoked = False def revoke self - None: """Mark the session credentials as revoked. In the current single-process model this is bookkeeping; in a multi-process / server scenario it would also invalidate the token in a shared store. """ self. revoked = True Rotate: generate a new token so any stale reference is useless. self.session token = secrets.token urlsafe 32 def repr self - str: Never include the token itself in repr/log output. return f"SessionCredentials id={self.session id r}, revoked={self. revoked} " The token is injected into the container env as AGENT SESSION TOKEN by the harness. Any harness-authenticated tool e.g. a future github api tool uses this token rather than a long-lived host credential. revoke is called in the finally block of main on session end. Recreating the container rotates the credentials, which DockerSandbox already does once per session. Security is great. But you also have to think about what happens when your security measures break down. For the last measure of security, we will ask ourselves: if something goes wrong anyway, can we see it, and can we stop it? AuditLog writes one JSON object per line to logs/audit-