Skip to content

Update pipeline.py #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Dayxxn343 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update pipeline.py #9

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 23 additions & 65 deletions viper/core/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,12 @@
from viper.generator.context.vuln_type_classifier import VulnTypeClassifier
from viper.generator.context.function_selector import FunctionSelector
from viper.generator.context.snippet_selector import SnippetSelector
from viper.llm.ollama_client import OllamaLLMClient
from viper.generator.llm.ollama_client import OllamaLLMClient

from viper.generator.poc_generator import (
PoCGenerationContext,
PoCGenerator,
)
from viper.reporter.analyzer import analyze_vulnerability
from viper.reporter.builder import build_vex_document
from viper.reporter.exporter import export_vex_json


class Pipeline:
Expand All @@ -38,13 +35,9 @@ def __init__(self, repo: str):
self.sbom_path = None
self.vuln_path = None
self.codeql_result_path = None
self.report_path = None

self.sbom_result = {}
self.vulnerabilities = []
self.analysis_contexts = []
self.poc_results = []
self.report_results = []

def _banner(self, title: str) -> str:
width = max(len(title) + 8, 52)
Expand Down Expand Up @@ -103,10 +96,7 @@ def analyze(self) -> None:
self._detail(f"Target repository: {self.repo}")

self._module_start("REPOSITORY PREPARATION")
self.repo_path = RepoCloner().prepare(
self.repo,
install_dependencies=False,
)
self.repo_path = RepoCloner().prepare(self.repo)
self._module_end("REPOSITORY PREPARATION", str(self.repo_path))
self._detail(f"Prepared repository: {self.repo_path}")

Expand All @@ -117,7 +107,6 @@ def analyze(self) -> None:
)
self._module_end("SBOM GENERATION", str(self.sbom_path))
self._detail(f"SBOM saved to: {self.sbom_path}")
self.sbom_result = json.loads(Path(self.sbom_path).read_text(encoding="utf-8"))

self._module_start("VULNERABILITY SCAN")
self.vuln_path = GrypeRunner().run(
Expand Down Expand Up @@ -232,7 +221,6 @@ def generate_poc(self) -> None:
)

generator = PoCGenerator(llm_client=llm_client)
self.poc_results = []

for context in self.analysis_contexts:
cve_id = context["cve_id"]
Expand Down Expand Up @@ -271,31 +259,12 @@ def generate_poc(self) -> None:
output_dir = Path("outputs/pocs") / cve_id
saved_path = generator.save(generated, output_dir=output_dir)

self.poc_results.append(
{
"cve_id": cve_id,
"success": True,
"verified": False,
"log": generated.explanation,
"poc_path": str(saved_path),
}
)

self._detail(f"Generated PoC path: {saved_path}")
self._detail(f"Explanation: {generated.explanation}")

except Exception as e:
self._console_spaced(f"failed for {cve_id}", color="red", level="error")
self._detail(f"PoC generation failed for {cve_id}: {e}")
self.poc_results.append(
{
"cve_id": cve_id,
"success": False,
"verified": False,
"log": str(e),
"poc_path": None,
}
)

def _load_vulnerabilities(self, vuln_path: str | Path) -> list[dict]:
vuln_path = Path(vuln_path)
Expand Down Expand Up @@ -337,49 +306,38 @@ def _load_vulnerabilities(self, vuln_path: str | Path) -> list[dict]:
def validate_poc(self) -> None:
stage = "VALIDATE POC"
with self._stage_indicator(stage):
self._detail(f"Target repository: {self.repo}")

# TODO

def report(self) -> None:
stage = "REPORT"
with self._stage_indicator(stage):
self._detail(f"Target repository: {self.repo}")

if not self.analysis_contexts:
self._detail("No analysis context found")
return

analysis_results = []
from viper.validator.poc_runner import PoCRunner
runner = PoCRunner()

self.validation_results = []

for context in self.analysis_contexts:
cve_id = context["cve_id"]
package_name = context["package_name"]
package_version = context["package_version"] or ""

poc_result = next(
(result for result in self.poc_results if result["cve_id"] == cve_id),
{},
)
poc_dir = Path("outputs/pocs") / cve_id

analysis_result = analyze_vulnerability(
self._module_start(f"VALIDATING {cve_id}")

# 가상 환경에서 PoC 동적 실행 및 결과 수집
result = runner.run_validation(
cve_id=cve_id,
package_name=package_name,
package_version=package_version,
sbom_result=self.sbom_result,
poc_result=poc_result,
repo_path=self.repo_path,
poc_dir=poc_dir
)
analysis_results.append(analysis_result)

self.validation_results.append(result)
self._module_end(f"VALIDATING {cve_id}", result["status"])
self._detail(f"[{cve_id}] Result Details: {result}")

vex_document = build_vex_document(analysis_results)
self.report_path = export_vex_json(
vex_document=vex_document,
output_path="outputs/vex/vex.json",
)
self.report_results = analysis_results
def report(self) -> None:
stage = "REPORT"
with self._stage_indicator(stage):
self._detail(f"Target repository: {self.repo}")

self._detail(f"VEX report saved to: {self.report_path}")
self._detail(f"VEX statement count: {len(analysis_results)}")
# TODO

def run(self) -> None:
self._console(self._logo(), color="blue")
Expand All @@ -394,4 +352,4 @@ def run(self) -> None:
self._detail(f"Full pipeline failed: {e}")
return

self._console("VIPER pipeline completed", color="green")
self._console("VIPER pipeline completed", color="green")