Skip to main content
YAK

Code Audit: AI-Powered Extensions to IRify and Automated PoC Generation

· 11 min read
Yak ProjectYak Project

Quick Overview

  • Applicable scenarios: Security teams that need to perform phased code audits on large projects, organize vulnerability evidence, and get assisted SyntaxFlow rule generation.
  • Implementation approach: First explore the project and build a scan plan, then audit by category, verify each piece of evidence one by one, generate a report, and produce PoC examples or rule suggestions driven by risk.
  • Key conclusion: The value of AI lies in organizing exploration, verification, and reporting into a traceable workflow; high-risk conclusions and PoCs still need to be reviewed in isolated, authorized environments.

FAQ

Does AI audit jump straight to vulnerability conclusions?

No. The workflow described breaks project exploration, scan planning, categorized audit, evidence recording, and report generation into independent phases, requiring conclusions to trace back to specific evidence.

Can auto-generated PoCs be used directly against production targets?

No. PoCs are intended for risk reproduction and verification within an authorized scope; their inputs, impact, and execution boundaries should be reviewed in an isolated environment.

As Memfit AI's new features expand and the related mechanisms mature, we can now extend existing AI capabilities to our other products. This article focuses on the new attempts we have made in the AI-powered extension of the IRify tool.

Pure AI Code Audit

In Yaklang's AI ReAct system, "code security audit" is registered as an independent loop type. Once triggered by the user, the runtime first constructs an audit state object that persists throughout the entire process and records the AI working directory. All staged artifacts converge into an audit subdirectory under that directory, preventing them from mingling with other tasks.

Currently there are two entry points in IRify:

  1. In the lower-right corner of project compilation, select: Perform AI Audit (the frontend needs an update; this will be fixed in a later version)

  1. On the AI Agent page, select Focus Mode: Code Security Audit

The audit state spans four phases, recording the project name and path, the tech stack and entry-point summary, the reconnaissance report path and note file list, the structured finding list from the scan phase, observations and uncertain leads across each scan category, confirmed and excluded conclusions from the verification phase, and the final report content and path.

Next, let's test using Yaklang's built-in vulnerability lab.

Phase 1: Project Exploration

The main orchestrator does not perform the directory walkthrough itself. Instead, it creates a "directory exploration" sub-loop, passes the expected output path of the reconnaissance notes and the exploration working directory into it via variables, and then runs that loop as an independent subtask.

During exploration, the program first reads go.mod to load dependencies, then uses various built-in tools to obtain the project structure.

After exploration finishes, the main orchestrator reclaims the results from the sub-loop's named variables: the real project path and name, the tech-stack description, the entry-point summary, the aggregated reconnaissance report file path, and the list of note paths produced by writing files during exploration, then writes them back into the shared audit state.

Phase 2: Scan Plan Plus Category-by-Category Serial Two-Stage Audit

Phase 2 itself is split into two layers: the outer category orchestration and the inner ReAct scan for a single category.

The outer layer first uses a single "fast" model call to let the AI decide, within structured output, which default vulnerability categories to scan this round and whether the user has additionally specified custom categories (with keywords and supplementary notes). If this planning fails, it falls back to the full built-in default category set, ensuring the workflow is not interrupted.

Each vulnerability category adopts a two-stage strategy:

  • The first stage is the keyword-search stage: based on the Sink semantic hints for that category and the tech stack, the model decides its own grep strategy, emphasizing the "only return matching files" mode to avoid truncated or missing results; after each search round, it merges the absolute paths of hit files into the pending-audit list via a dedicated action, and may continue with multiple search rounds.
  • The second stage is the file-by-file audit stage: it reads the contents of each file locked in by the previous stage in turn; if a problem is confirmed, it submits a structured finding. There is a lower bound on confidence, and submissions below it are rejected to curb casually reporting vulnerabilities.

The screenshot above shows the first stage of the keyword search. The program first found high-risk files that obviously contain vulnerability examples, such as vul_cmdi.go, and then enters the second stage to search line by line.

It reads potentially problematic code snippets via tools and hands them to the AI to complete the audit.

Phase 3: Item-by-Item Verification and Evidence Recording

The dedicated actions in the verification phase mainly fall into several categories:

  1. Trace tracking nodes step by step along the data flow, recording the step, variable, location, inferred source type, and so on into the timeline.

  2. Record the filtering or validation considered to be present, along with its effectiveness judgment.

  3. Provide a final conclusion for a finding, selecting from only three states — confirmed, excluded, or needs manual confirmation — along with confidence and rationale, and optionally refine the data flow, exploit method, and remediation suggestions.

  4. Read the reconnaissance background report again.

  5. Once every finding has a conclusion, invoke complete-verification and exit this phase.

Phase 4: Report Generation and Fallback

The system creates a "report generation" sub-loop, injecting the final Markdown report path, the writing task description, and the list of reference files that must be read into the sub-loop variables. A loop specialized in writing long-form content reads these inputs and writes them to disk.

Risk-Based Code Audit

In addition to the pure-AI code audit, IRify also supports code auditing driven by existing risk information. For this precise yet lightweight feature, we chose to implement it with forge.

For example, suppose we have already scanned a project and obtained the following risk information:

Now, to perform an AI audit on the high-risk item number 1, "Detected Golang server-side template injection vulnerability," we can search for ssa-risk in the AI Agent's tool library to locate the target. After clicking it, enter the risk number: 1 in the popup.

After execution, the AI returns a concise vulnerability analysis report.

Risk-Driven Intelligent PoC Example Generation

Currently IRify has built-in code constructors for two languages — Python and Yak — that can produce correct code free of syntax errors. Combined with the risk-based code audit feature (ssa-risk in the tool library), it can directly generate PoC examples from risks.

Given existing static-analysis risk results, multi-step reasoning turns "risk metadata + source/fragment context + vulnerability type" into "a proof-of-concept script you can take into an authorized environment for verification." In the skill library, search for "Intelligently Generate PoC from SSA Risk". After clicking it, configure the risk serial number to audit in the popup, and PoC example generation begins.

The program first runs ssa-risk once to audit the risk-related content, and then reasons from the following directions:

  1. Determine whether the source code can be followed

Check whether the associated project is in-memory compiled and whether source is available: if followable, continue using listing/grep/read to fill in routes and call chains; if not, mainly rely on the code snippets that come with the Risk.

  1. Fill in context (optional)

When source is readable, list directories on demand, search routes/configs, and read key files to extract the "real entry points, methods, parameters, and authentication."

  1. Classification + attack-surface analysis

Categorize the risk into specific exploit ideas like SQLi/XSS/SSRF, align the data flow, and set the PoC technical route (this step focuses on analysis and does not require writing a long essay first).

  1. Find a PoC template

Use the template-retrieval tool to fetch Python-side skeletons for request construction, result parsing, and so on by vulnerability type, reducing AI hallucination.

The images below show how the program determines whether the target project is reachable and performs data-flow analysis:

This system supports configuring the target project's real runtime environment address (e.g., http://your-app:8080) in the target_base_url parameter. The system will then automatically run an intelligent security-analysis-and-verification workflow and ultimately generate a report and PoC example code:

SyntaxFlow Rule Beautification

"Beautification" here mainly means reformatting a .sf rule text according to a unified standard for output: indentation, line breaks, and the structure of desc / alert and the rule body become more consistent; when options are supplied, it can also complete or normalize metadata (such as the rule ID, fields required by the description, etc.).

Currently SyntaxFlow rule beautification has two entry points: the upper-right corner when creating a new rule, and the upper-right corner of the rule-writing page in the code audit.

Note: before beautification begins, make sure the current rule has no syntax errors, because we run a SyntaxFlow syntax check both immediately before and after beautification to ensure beautification itself does not break the rule's syntax. Furthermore, rule beautification uses a strict metadata-replacement logic, so there's no need to worry about it damaging the rule's own content.

SyntaxFlow Rule Auto-Generation

This feature performs retrieval over the built-in rule-example knowledge pack, then combines the single-file editing toolchain with SyntaxFlow compilation checks to iteratively write a valid .sf file.

On the AI Agent page, select Focus Mode: write_syntaxflow_rule

System behavior: extract samples (language / virtual filename / code body) and save them → when generating the rule, embed the samples as an UNSAFE positive-test block at the end of the rule → after generation, run a syntax/sample-hit self-check; if it fails, iteratively modify the rule until it passes.

  1. Knowledge retrieval
  • Archive retrieval: search syntaxflow-aikb by pattern for hit fragments (with context).

  • Vector retrieval (enabled if available): query syntaxflow-aikb-rag with multiple complete sentences; deduplicate the results and rerank them (leaning toward content with executable-rule features and de-emphasizing pure long-form prose).

  1. Knowledge compression
  • The raw results assembled from grep + RAG are often large; they are first merged into one large block of text, and then a compression call tailored to the current need and search conditions is made (targeting roughly the 10KB scale in the implementation).
  1. Persist and pin the file
  • If there are vulnerability samples, they are first written as temporary sample files; then either open the user-specified .sf or create a new empty .sf draft, and pin that path as the subsequent editing target.
  1. Multi-round write-edit + verification
  • The model maintains the rule body via actions such as writing/modifying the rule; each change goes through SyntaxFlow syntax and format diagnostics (errors block further indiscriminate edits). If samples are supplied, it also runs a positive-sample self-check via check-syntaxflow-syntax until the rule hits as expected on the samples.

To use it, just paste a snippet of vulnerability example code and it will automatically complete the SyntaxFlow rule, run syntax detection and usability checks, and finally score the rule.

You can find the generated SyntaxFlow rule in the file tree.

Future Extensions

The existing AI code audit still has many difficult code-level problems unsolved. Affected by model quality and context limits, AI audits perform poorly on large projects; the core issue is that stability is too low, and multiple audits rarely produce the same result. We are trying to use the stable risk artifacts from code scanning to shore up the unstable pure-AI code audit.

IRify represents the trend of code security audit moving toward intelligence, automation, and precision. Through Yaklang's AI ReAct system, it combines human experts' experiential patterns with AI's fast analysis capabilities, building the core competency of a new-generation security-analysis platform. The system will continue to optimize its performance in scenarios such as complex projects, multi-language support, and cloud-native environments, providing more comprehensive intelligent support for the software security development lifecycle.


This article was first published on the Yak Project official account. Read the original.