feat: add taskflow prompts to handle libraries in addition to applications

scripts/audit/run_audit.sh

@@ -5,14 +5,19 @@
 set -e
 
 USE_ADVISORY=false
+MODEL_CONFIG_FLAG=""
 
 # Parse flags
-while [[ "$1" == --* ]]; do
+while [[ "$1" == -* ]]; do
   case "$1" in
     --advisory)
       USE_ADVISORY=true
       shift
       ;;
+    -m)
+      MODEL_CONFIG_FLAG="-m $2"
+      shift 2
+      ;;
     *)
       echo "Unknown option: $1"
       exit 1
@@ -21,20 +26,21 @@ while [[ "$1" == --* ]]; do
 done
 
 if [ -z "$1" ]; then 
-  echo "Usage: $0 [--advisory] <repo>"; 
+  echo "Usage: $0 [--advisory] [-m model_config] <repo>"; 
   exit 1; 
 fi
 
-python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.fetch_source_code -g repo="$1"
-python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.identify_applications -g repo="$1"
-python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.gather_web_entry_point_info -g repo="$1"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.fetch_source_code -g repo="$1"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.identify_applications -g repo="$1"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.gather_web_entry_point_info -g repo="$1"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.gather_security_entry_point_info -g repo="$1"
 
 if [ "$USE_ADVISORY" = true ]; then
-  python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.fetch_security_advisories -g repo="$1"
+  python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.fetch_security_advisories -g repo="$1"
 fi
 
-python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.classify_application_local -g repo="$1" -g use_advisory="$USE_ADVISORY"
-python -m seclab_taskflow_agent -t seclab_taskflows.taskflows.audit.audit_issue_local_iter -g repo="$1" -g use_advisory="$USE_ADVISORY"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.classify_application_local -g repo="$1" -g use_advisory="$USE_ADVISORY"
+python -m seclab_taskflow_agent $MODEL_CONFIG_FLAG -t seclab_taskflows.taskflows.audit.audit_issue_local_iter -g repo="$1" -g use_advisory="$USE_ADVISORY"
 
 set +e
 

src/seclab_taskflows/mcp_servers/local_gh_resources.py

@@ -71,6 +71,18 @@ async def _fetch_file(url, headers, params):
     return await _fetch_file(url, headers=headers, params=params)
 
 
+def _parse_content_disposition_filename(header: str | None) -> str | None:
+    """Extract the filename from a Content-Disposition header value."""
+    if not header:
+        return None
+    for raw_part in header.split(";"):
+        part = raw_part.strip()
+        if part.lower().startswith("filename="):
+            filename = part.split("=", 1)[1].strip().strip('"')
+            return filename
+    return None
+
+
 async def _fetch_source_zip(owner: str, repo: str, tmp_dir):
     """Fetch the source code."""
     url = f"https://api.github.com/repos/{owner}/{repo}/zipball"
@@ -83,6 +95,8 @@ async def _fetch_source_zip(owner: str, repo: str, tmp_dir):
         async with httpx.AsyncClient() as client:
             async with client.stream("GET", url, headers=headers, follow_redirects=True) as response:
                 response.raise_for_status()
+                content_disposition = response.headers.get("content-disposition")
+                source_filename = _parse_content_disposition_filename(content_disposition)
                 expected_path = Path(tmp_dir) / owner / f"{repo}.zip"
                 resolved_path = expected_path.resolve()
                 if os.path.commonpath([resolved_path, Path(tmp_dir).resolve()]) != str(Path(tmp_dir).resolve()):
@@ -92,6 +106,11 @@ async def _fetch_source_zip(owner: str, repo: str, tmp_dir):
                 async with aiofiles.open(f"{tmp_dir}/{owner}/{repo}.zip", "wb") as f:
                     async for chunk in response.aiter_bytes():
                         await f.write(chunk)
+        metadata = {"source_filename": source_filename}
+        metadata_path = Path(tmp_dir) / owner / f"{repo}_source_metadata.json"
+        async with aiofiles.open(metadata_path, "w") as f:
+            await f.write(json.dumps(metadata, indent=2))
+
         return f"source code for {repo} fetched successfully."
     except httpx.RequestError as e:
         return f"Error: Request error: {e}"
@@ -113,6 +132,7 @@ async def fetch_repo_from_gh(owner: str, repo: str):
     source_path = Path(f"{LOCAL_GH_DIR}/{owner}/{repo}.zip")
     if not source_path.exists():
         return result
+
     return f"Downloaded source code to {owner}/{repo}.zip"