main.py

import argparse
import json
import os
from typing import Dict, List
import logging

from python_api_extractor.extracter import api_extracter
from python_ast_visualizer.utils import ast_to_png
from search_source.container_extracter import layer_extracter
from trivy_extracter.trivy_module import trivy_func

TRIVY_OUTPUT_PATH = "DB/trivy_analysis_result.json"
LIB2CVE2API_PATH = "DB/LIB2CVE2API.json"
AST_RESULT_PATH = "DB/anaysis_result"

def load_api_list(json_path: str) -> list:
    with open(json_path, encoding='utf-8') as f:
        data = json.load(f)

    api_list = []
    for pkg_versions in data.values():
        for entry in pkg_versions.values():
            # entry['apis']는 {module: [func1, func2, ...], ...}
            for funcs in entry.get('apis', {}).values():
                api_list.extend(funcs)
    return sorted(set(api_list))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("image_tar", help="Path to Docker image tar file")

    args = parser.parse_args()


    # extract container image info by Trivy
    logging.info("step 1. checksec with trivy")
    trivy_func.scan_vulnerabilities(args.image_tar, TRIVY_OUTPUT_PATH )
    logging.info("step 1 Done!")


    # extract API set from "DB/trivy_analysis_result.json"
    logging.info("step 2. Map API set with lib-cve data")

    with open(TRIVY_OUTPUT_PATH, encoding="utf-8") as f:
        data = json.load(f)

    combined: Dict[str, Dict[str, Dict[str, List[str]]]] = api_extracter.build_cve_api_mapping(data)
    output_str = json.dumps(combined, indent=2, ensure_ascii=False)

    with open(LIB2CVE2API_PATH, 'w', encoding='utf-8') as out_f:
            out_f.write(output_str)

    logging.info("step 2 Done!")


    logging.info("step 3. Search source file from container and extract")
    source_path = os.path.join("DB" + os.path.splitext(os.path.basename(args.image_tar))[0] + "source")
    layer_extracter.extract_app_layer(
        args.image_tar,
        output_dir = source_path,
        auto_detect=True,
        include_filter=".py"
    )
    logging.info("step 3 Done!")


    # Step 4: AST-based call flow analysis
    logging.info("step 4. AST analysis")

    container_base = os.path.splitext(os.path.basename(args.image_tar))[0]
    analysis_dir = os.path.abspath(os.path.join("DB", f"{container_base}_analysis_result"))
    os.makedirs(analysis_dir, exist_ok=True)
    target_list = load_api_list(LIB2CVE2API_PATH)
    force = len(target_list) == 0
    targets = ast_to_png.parse_target_calls(target_list)

    files = []
    if os.path.isdir(source_path):
        for root, _, fns in os.walk(source_path):
            for fn in fns:
                if fn.endswith('.py'):
                    files.append(os.path.join(root, fn))
        logging.info(f"Collected {len(files)} Python files for analysis from {source_path}")
        base = source_path.rstrip(os.sep)
    else:
        files = [source_path]
        base = os.path.dirname(source_path) or '.'
        logging.info(f"Single file mode: {source_path}")

    image_prefix = os.path.join(analysis_dir, "ast_graph")
    external_apis, internal_only_apis, unused_apis = ast_to_png.visualize_call_flow(
        files, base, image_prefix, targets, force
    )
    results = {
        "external": external_apis,
        "internal": internal_only_apis,
        "unused": unused_apis
    }
    json_path = os.path.join(analysis_dir, "api_info.json")
    with open(json_path, 'a', encoding='utf-8') as jf:
        json.dump(results, jf, indent=2, ensure_ascii=False)
    logging.info(f"After JSON write, contents of {analysis_dir}: {os.listdir(analysis_dir)}")

    logging.info(f"AST analysis results and images written to {analysis_dir}")