ai_ops/src/security/schemas.ts

import { z } from "zod";

function dedupe(values: readonly string[]): string[] {
  const out: string[] = [];
  const seen = new Set<string>();

  for (const value of values) {
    if (seen.has(value)) {
      continue;
    }
    seen.add(value);
    out.push(value);
  }

  return out;
}

const stringTokenSchema = z
  .string()
  .trim()
  .min(1)
  .refine((value) => !value.includes("\0"), "String token cannot include null bytes.");

export const toolClearancePolicySchema = z
  .object({
    allowlist: z.array(stringTokenSchema).default([]),
    banlist: z.array(stringTokenSchema).default([]),
  })
  .strict();

export type ToolClearancePolicy = z.infer<typeof toolClearancePolicySchema>;

export function parseToolClearancePolicy(input: unknown): ToolClearancePolicy {
  const parsed = toolClearancePolicySchema.parse(input);
  return {
    allowlist: dedupe(parsed.allowlist),
    banlist: dedupe(parsed.banlist),
  };
}

export const shellValidationPolicySchema = z
  .object({
    allowedBinaries: z.array(stringTokenSchema).min(1),
    worktreeRoot: stringTokenSchema,
    protectedPaths: z.array(stringTokenSchema).default([]),
    requireCwdWithinWorktree: z.boolean().default(true),
    rejectRelativePathTraversal: z.boolean().default(true),
    enforcePathBoundaryOnArguments: z.boolean().default(true),
    allowedEnvAssignments: z.array(stringTokenSchema).default([]),
    blockedEnvAssignments: z.array(stringTokenSchema).default([]),
  })
  .strict();

export type ShellValidationPolicy = z.infer<typeof shellValidationPolicySchema>;

export function parseShellValidationPolicy(input: unknown): ShellValidationPolicy {
  const parsed = shellValidationPolicySchema.parse(input);
  return {
    ...parsed,
    allowedBinaries: dedupe(parsed.allowedBinaries),
    protectedPaths: dedupe(parsed.protectedPaths),
    allowedEnvAssignments: dedupe(parsed.allowedEnvAssignments),
    blockedEnvAssignments: dedupe(parsed.blockedEnvAssignments),
  };
}

export const executionEnvPolicySchema = z
  .object({
    inherit: z.array(stringTokenSchema).default(["PATH", "HOME", "TMPDIR", "TMP", "TEMP", "LANG", "LC_ALL"]),
    scrub: z.array(stringTokenSchema).default([]),
    inject: z.record(z.string(), z.string()).default({}),
  })
  .strict();

export type ExecutionEnvPolicy = z.infer<typeof executionEnvPolicySchema>;

export function parseExecutionEnvPolicy(input: unknown): ExecutionEnvPolicy {
  const parsed = executionEnvPolicySchema.parse(input);
  return {
    inherit: dedupe(parsed.inherit),
    scrub: dedupe(parsed.scrub),
    inject: { ...parsed.inject },
  };
}

export const commandTargetSchema = z
  .object({
    binary: stringTokenSchema,
    args: z.array(stringTokenSchema).default([]),
  })
  .strict();

export type CommandTarget = z.infer<typeof commandTargetSchema>;

export const commandTargetsSchema = z.array(commandTargetSchema);

export function parseCommandTargets(input: unknown): CommandTarget[] {
  const parsed = commandTargetsSchema.parse(input);
  return parsed.map((target) => ({
    binary: target.binary,
    args: [...target.args],
  }));
}

export const parsedShellAssignmentSchema = z
  .object({
    raw: stringTokenSchema,
    key: stringTokenSchema,
    value: z.string(),
  })
  .strict();

export type ParsedShellAssignment = z.infer<typeof parsedShellAssignmentSchema>;

export const parsedShellCommandSchema = z
  .object({
    binary: stringTokenSchema,
    args: z.array(stringTokenSchema).default([]),
    flags: z.array(stringTokenSchema).default([]),
    assignments: z.array(parsedShellAssignmentSchema).default([]),
    redirects: z.array(stringTokenSchema).default([]),
    words: z.array(stringTokenSchema).default([]),
  })
  .strict();

export type ParsedShellCommand = z.infer<typeof parsedShellCommandSchema>;

export const parsedShellScriptSchema = z
  .object({
    commandCount: z.number().int().nonnegative(),
    commands: z.array(parsedShellCommandSchema).default([]),
  })
  .strict()
  .refine((value) => value.commandCount === value.commands.length, {
    message: "commandCount must match commands.length",
    path: ["commandCount"],
  });

export type ParsedShellScript = z.infer<typeof parsedShellScriptSchema>;

export function parseParsedShellScript(input: unknown): ParsedShellScript {
  const parsed = parsedShellScriptSchema.parse(input);
  return {
    commandCount: parsed.commandCount,
    commands: parsed.commands.map((command) => ({
      binary: command.binary,
      args: [...command.args],
      flags: [...command.flags],
      assignments: command.assignments.map((assignment) => ({
        raw: assignment.raw,
        key: assignment.key,
        value: assignment.value,
      })),
      redirects: [...command.redirects],
      words: [...command.words],
    })),
  };
}

export type SecurityViolationHandling =
  | "hard_abort"
  | "validation_fail"
  | "dangerous_warn_only";

export const securityViolationHandlingSchema = z.union([
  z.literal("hard_abort"),
  z.literal("validation_fail"),
  z.literal("dangerous_warn_only"),
]);

export function parseSecurityViolationHandling(input: unknown): SecurityViolationHandling {
  return securityViolationHandlingSchema.parse(input);
}