orchestrator.py

import argparse
import subprocess
import json
import os
import shutil
import pathlib
import csv
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel


# ==========================
# Load fine‑tuned model
# ==========================
base_model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, "./../deepseek-lora-v5/checkpoint-9702")
model.eval()

system_prompt = """Here is a JavaScript AWS Lambda function, along with its least privilege IAM policy and trigger configuration.
Based on the given IAM permissions and trigger types, insert a security monitoring logic into the function to enforce least privilege access.
Only return the modified function code — no explanation or comments.
"""

def call_your_llm(payload: dict):
    input_obj = {
        "code": payload["source_code"],
        "least_privilege_policy": payload["least_privilege_policy"],
        "triggers": payload["triggers"]
    }
    print(input_obj)
    user_prompt = json.dumps(input_obj, indent=2)
    messages = [
        { 'role': 'system', 'content': system_prompt },
        { 'role': 'user', 'content': user_prompt }
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)
    outputs = model.generate(
        inputs,
        max_new_tokens=2048,
        do_sample=False,
        top_k=50,
        top_p=0.95,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id
    )
    result = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
    try:
        return json.loads(result)
    except:
        return {"modified_code": result, "iam_policy": {}}

# ==========================
# Parse CLI arguments
# ==========================
parser = argparse.ArgumentParser()
parser.add_argument("--source-file", required=True, help="Path to the source code file to analyze")
parser.add_argument("--triggers", nargs="+", required=True,
                    help="Trigger list (must be one or more of: event, api, direct)")
parser.add_argument("--account-id", required=True, help="AWS Account ID")
parser.add_argument("--region", required=True, help="AWS Region")
args = parser.parse_args()

# Validate triggers
valid_triggers = {"event", "api", "direct"}
for t in args.triggers:
    if t not in valid_triggers:
        raise ValueError(f"Invalid trigger: {t}. Only event, api, direct are allowed.")

print(f"[INFO] Triggers: {args.triggers}")

# Detect language by extension
ext = os.path.splitext(args.source_file)[1].lower()
if ext == ".js":
    language = "javascript"
elif ext == ".py":
    language = "python"
elif ext == ".go":
    language = "go"
else:
    raise ValueError(f"Unsupported file extension: {ext}. Only js, py, go are supported.")

print(f"[INFO] Detected language: {language}")

# Create CodeQL database
source_root = os.path.dirname(os.path.abspath(args.source_file))
source_filename = pathlib.Path(args.source_file).stem
db_path = f"./temp_db_{source_filename}"

if os.path.exists(db_path):
    shutil.rmtree(db_path)

print(f"[INFO] Creating CodeQL database from source root: {source_root}")
print(f"[INFO] DB path: {db_path}")
subprocess.run([
    "codeql", "database", "create", db_path,
    f"--language={language}",
    f"--source-root={source_root}"
], check=True)

# Run CodeQL analysis with CSV output
OUTPUT_DIR = "output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
OUTPUT_CSV = os.path.join(OUTPUT_DIR, "codeql_out.csv")
QUERY_FILE = "queries/js.ql"
BQRS_FILE = os.path.join(OUTPUT_DIR, "codeql_out.bqrs")

print(f"[INFO] Running CodeQL analysis on DB: {db_path}")
# subprocess.run([
#     "codeql", "database", "analyze", db_path,
#     QUERY_FILE,
#     "--format=csv",
#     f"--output={OUTPUT_CSV}"
# ], check=True)
# Run query to produce BQRS
# Run query and produce BQRS in output directory
import time

start = time.perf_counter()
subprocess.run([
    "codeql", "query", "run",
    QUERY_FILE,
    f"--database={db_path}",
    f"--output={BQRS_FILE}"
], check=True)


# Decode BQRS to CSV
subprocess.run([
    "codeql", "bqrs", "decode",
    "--format=csv",
    f"--output={OUTPUT_CSV}",
    BQRS_FILE
], check=True)

mid = time.perf_counter()


# Read CodeQL results from CSV
analysis_result = []
with open(OUTPUT_CSV, newline='') as csvfile:
    reader = csv.reader(csvfile)
    headers = next(reader, None)  # skip header if present
    for row in reader:
        analysis_result.append(row)

# Read source code
with open(args.source_file) as f:
    source_code = f.read()

# ----------------------------
# Build least_privilege_policy from CSV
# ----------------------------
least_privilege_policy = [
    f"{row[0]}:{row[1]}"
    for row in analysis_result
]


# --- CSV 읽은 후 JSON 결과 생성 ---
least_privilege_policy2 = []
for row in analysis_result:
    permission = row[0]
    resource_value = row[1]
    service, action = permission.split(":", 1)

    if service == "dynamodb":
        arn = f"arn:aws:dynamodb:{args.region}:{args.account_id}:table/{resource_value}"
    elif service == "lambda":
        arn = f"arn:aws:lambda:{args.region}:{args.account_id}:function:{resource_value}"
    elif service == "s3":
        arn = f"arn:aws:s3:::{resource_value}"
    elif service == "sns":
        arn = f"arn:aws:sns:{args.region}:{args.account_id}:{resource_value}"
    elif service == "sqs":
        arn = f"arn:aws:sqs:{args.region}:{args.account_id}:{resource_value}"
    else:
        # 기본 포맷 (서비스별 세부 포맷 필요시 여기에 추가)
        arn = f"arn:aws:{service}:{args.region}:{args.account_id}:{resource_value}"


    least_privilege_policy2.append({
        "action": permission,
        "resource": arn
    })

# JSON 파일로도 저장
with open(os.path.join(OUTPUT_DIR, "codeql_result.json"), "w") as f:
    json.dump(least_privilege_policy2, f, indent=2)


# ----------------------------
# Call LLM with proper structure
# ----------------------------
response = call_your_llm({
    "source_code": source_code,
    "least_privilege_policy": least_privilege_policy,
    "triggers": args.triggers
})
end = time.perf_counter()

# Save results
with open(os.path.join(OUTPUT_DIR, "modified_lambda.out"), "w") as f:
    f.write(response["modified_code"])
with open(os.path.join(OUTPUT_DIR, "iam_policy.json"), "w") as f:
    json.dump(response["iam_policy"], f, indent=2)

print(torch.cuda.is_available())

print("[DONE] Results saved to output/ directory.")
print(f"CodeQL 실행 시간: {mid - start:.3f}초")
print(f"LLM 호출 시간: {end - mid:.3f}초")
print(f"총 처리 시간: {end - start:.3f}초")