Spaces:
Running on A10G
Running on A10G
File size: 5,232 Bytes
95cbc5b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import argparse
import json
import subprocess
import sys
from dataclasses import asdict
from pathlib import Path
from .scanner import CommitGuardScanner
def cmd_scan(args):
diff_text = ""
if getattr(args, "diff", None):
if args.diff in ("-", "/dev/stdin"):
diff_text = sys.stdin.read()
else:
diff_text = Path(args.diff).read_text(encoding="utf-8")
elif getattr(args, "staged", False):
diff_text = subprocess.check_output(["git", "diff", "--staged"], text=True)
elif getattr(args, "commit", None):
diff_text = subprocess.check_output(["git", "show", args.commit], text=True)
elif getattr(args, "pr", None):
diff_text = subprocess.check_output(["gh", "pr", "diff", args.pr], text=True)
else:
print("Must specify one of --diff, --staged, --commit, or --pr")
sys.exit(1)
if not diff_text.strip():
print("No diff found to scan.")
sys.exit(0)
print(f"Loading model ({args.model})...", file=sys.stderr)
scanner = CommitGuardScanner(model_path=args.model, is_lora=args.is_lora, base_model=args.base_model)
print(f"Scanning diff ({len(diff_text)} chars)...", file=sys.stderr)
result = scanner.scan(diff_text)
if args.format == "json":
print(json.dumps(asdict(result), indent=2))
elif args.format == "text":
status = "VULNERABLE ⚠️" if result.is_vulnerable else "SAFE ✅"
print(f"\nVerdict: {status}")
if result.is_vulnerable:
print(f"CWE: {result.cwe}")
print(f"Exploit Sketch:\n {result.exploit_sketch}")
if result.parse_error:
print(f"\nParser Warning: {result.parse_error}")
elif args.format == "sarif":
# Minimal SARIF output stub
print("SARIF format not fully implemented yet.", file=sys.stderr)
print(json.dumps(asdict(result)))
if args.fail_on_vulnerable and result.is_vulnerable:
sys.exit(1)
def cmd_server(args):
from .server import main as server_main
server_main()
def cmd_eval(args):
# This is a bit hacky to reuse the script without modifying sys.path everywhere
# A cleaner approach would be moving evaluate.py into commitguard_env
REPO_ROOT = Path(__file__).resolve().parent.parent
eval_script = REPO_ROOT / "scripts" / "evaluate.py"
cmd = [sys.executable, str(eval_script)]
cmd.extend(args.eval_args)
subprocess.run(cmd, check=True)
def cmd_hook(args):
from .hooks import install_hook
if args.action == "install":
if args.pre_commit:
install_hook("pre-commit")
elif args.pre_push:
install_hook("pre-push")
else:
print("Please specify a hook type to install (e.g., --pre-commit or --pre-push)")
sys.exit(1)
def main():
parser = argparse.ArgumentParser(description="CommitGuard AI-paced security review")
subparsers = parser.add_subparsers(dest="command", required=True)
# 'scan' subcommand
scan_parser = subparsers.add_parser("scan", help="Scan a code diff for vulnerabilities")
source_group = scan_parser.add_mutually_exclusive_group(required=True)
source_group.add_argument("--diff", type=str, help="Path to a diff file")
source_group.add_argument("--staged", action="store_true", help="Scan git staged changes")
source_group.add_argument("--commit", type=str, help="Scan a specific git commit (e.g., HEAD)")
source_group.add_argument("--pr", type=str, help="Scan a GitHub PR URL or ID (requires gh cli)")
scan_parser.add_argument("--model", type=str, default="inmodel-labs/commitguard-llama-3b", help="Model path or HF ID")
scan_parser.add_argument("--base-model", type=str, default=None, help="Base model if using LoRA")
scan_parser.add_argument("--is-lora", action="store_true", help="Whether the model is a LoRA adapter")
scan_parser.add_argument("--format", choices=["text", "json", "sarif"], default="text", help="Output format")
scan_parser.add_argument("--fail-on-vulnerable", action="store_true", help="Exit with code 1 if vulnerable")
# 'server' subcommand
server_parser = subparsers.add_parser("server", help="Start the OpenEnv environment server")
# server_main takes PORT from environment
# 'eval' subcommand
eval_parser = subparsers.add_parser("eval", help="Run the evaluation harness")
eval_parser.add_argument("eval_args", nargs=argparse.REMAINDER, help="Arguments passed to evaluate.py")
# 'hook' subcommand
hook_parser = subparsers.add_parser("hook", help="Manage git hooks")
hook_parser.add_argument("action", choices=["install"], help="Action to perform (e.g., install)")
hook_parser.add_argument("--pre-commit", action="store_true", help="Install pre-commit hook")
hook_parser.add_argument("--pre-push", action="store_true", help="Install pre-push hook")
args = parser.parse_args()
if args.command == "scan":
cmd_scan(args)
elif args.command == "server":
cmd_server(args)
elif args.command == "eval":
cmd_eval(args)
elif args.command == "hook":
cmd_hook(args)
if __name__ == "__main__":
main()
|