cwe_top25_filter.py

import os
import json
import openai
from datetime import datetime
import logging

CWE_TOP_25 = [
    "CWE-787", "CWE-79", "CWE-89", "CWE-416", "CWE-78", "CWE-20", "CWE-125", "CWE-22", "CWE-352", "CWE-434",
    "CWE-862", "CWE-476", "CWE-287", "CWE-190", "CWE-502", "CWE-77", "CWE-119", "CWE-798", "CWE-918", "CWE-306",
    "CWE-362", "CWE-269", "CWE-94", "CWE-863", "CWE-276"
]

def setup_logging():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_filename = f"cwe_top25_analysis_filtered_{timestamp}.log"
    
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    
    file_handler = logging.FileHandler(log_filename)
    file_handler.setLevel(logging.INFO)
    
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    return logger

def call_llm(messages):
    try:
        client = openai.OpenAI(
            api_key=os.environ["OPENAI_API_KEY"],
            base_url=os.environ["OPENAI_API_BASE"]
        )
        response = client.chat.completions.create(
            model="claude-3-5-sonnet",
            messages=messages
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        logging.error(f"Exception in call_llm: {str(e)}")
        return None

def parse_json(response):
    try:
        if response.find("```json") >= 0:
            start_pos = response.find("```json") + 7
            end_pos = response.rfind("```")
            response = response[start_pos:end_pos]
        else:
            start_pos = response.find("{")
            end_pos = response.rfind("}")
            response = response[start_pos:end_pos+1]
        return json.loads(response, strict=False)
    except Exception as e:
        print("Exception:", str(e))
        raise e

def analyze_cwe_top25(code_path):
    with open(code_path, 'r', encoding='utf-8') as file:
        code = file.read()
    
    cwe_list = ", ".join(CWE_TOP_25)
    prompt = f"""Analyze the following code for potential security risks based on the specific CWE Top 25 Most Dangerous Software Weaknesses list provided:

{code}

CWE Top 25 list to check against: {cwe_list}

Instructions:
1. Carefully examine the code for any violations of the CWE Top 25 list provided.
2. If you find any violations, list ALL of them. Do not limit the number of issues reported.
3. If you do not find any violations related to the CWE Top 25 list, return an empty list of risks.
4. For each identified risk, provide the following information:
   a. The CWE identifier (e.g., CWE-79)
   b. The name of the CWE
   c. A clear explanation of how the code violates this CWE
   d. A suggested mitigation strategy

Provide the output in the following JSON format:
{{
    "risks": [
        {{
            "cwe_id": "CWE-XXX",
            "cwe_name": "Name of the CWE",
            "explanation": "How the code violates this CWE",
            "mitigation": "Suggested mitigation strategy"
        }}
    ]
}}

Remember: If no violations are found, the "risks" array should be empty.
"""
    messages = [{"role": "user", "content": prompt}]
    return call_llm(messages)

def filter_false_positives(code_path, risks):
    with open(code_path, 'r', encoding='utf-8') as file:
        code = file.read()
    
    prompt = f"""Given the following code and a list of potential CWE Top 25 risks, please review each risk and determine if it's a true positive or a false positive. Only include true positives in the output.

Code:
{code}

Potential risks:
{json.dumps(risks, indent=2)}

Instructions:
1. Carefully review each risk in the context of the provided code.
2. Determine if each risk is a true positive (actually present in the code) or a false positive (incorrectly identified).
3. Only include true positives in the output.
4. Provide a brief explanation for each risk you include, stating why it's a true positive.

Provide the output in the following JSON format:
{{
    "filtered_risks": [
        {{
            "cwe_id": "CWE-XXX",
            "cwe_name": "Name of the CWE",
            "explanation": "Why this is a true positive",
            "mitigation": "Suggested mitigation strategy"
        }}
    ]
}}

If all risks are determined to be false positives, return an empty list for "filtered_risks".
"""
    messages = [{"role": "user", "content": prompt}]
    return call_llm(messages)

def get_all_files(directory):
    file_list = []
    for root, _, files in os.walk(directory):
        for file in files:
            file_list.append(os.path.relpath(os.path.join(root, file), directory))
    return file_list

def process_case(case_path, logger, file_count, total_cwe_issues):
    before_path = os.path.join(case_path, 'before')
    
    before_file = get_all_files(before_path)[0]
    before_code_path = os.path.join(before_path, before_file)

    logger.info(f"Processing case: {case_path}")

    logger.info("Analyzing code for CWE Top 25 risks...")
    analysis_result = analyze_cwe_top25(before_code_path)
    logger.info(f"Analysis result:\n{analysis_result}")

    parsed_result = parse_json(analysis_result)
    logger.info(f"Parsed JSON result:\n{json.dumps(parsed_result, indent=2)}")

    logger.info("Filtering false positives...")
    filtered_result = filter_false_positives(before_code_path, parsed_result['risks'])
    logger.info(f"Filtered result:\n{filtered_result}")

    filtered_parsed_result = parse_json(filtered_result)
    logger.info(f"Parsed filtered result:\n{json.dumps(filtered_parsed_result, indent=2)}")

    file_count += 1
    cwe_issues = len(filtered_parsed_result['filtered_risks'])
    total_cwe_issues += cwe_issues

    logger.info(f"Files analyzed: {file_count}, CWE Top 25 issues found in this file: {cwe_issues}, Total issues found so far: {total_cwe_issues}")

    return filtered_parsed_result, file_count, total_cwe_issues

def main():
    logger = setup_logging()
    logger.info("Starting CWE Top 25 analysis process with AI filtering")

    with open('small_file_cases.json', 'r') as f:
        cases = json.load(f)

    results = {}
    file_count = 0
    total_cwe_issues = 0
    
    for case in cases:
        try:
            case_result, file_count, total_cwe_issues = process_case(case, logger, file_count, total_cwe_issues)
            results[case] = case_result
            
            risks_count = len(case_result['filtered_risks'])
            if risks_count == 0:
                logger.info(f"Case {case}: No CWE Top 25 risks identified after filtering")
            else:
                logger.info(f"Case {case}: {risks_count} CWE Top 25 risks identified after filtering")
            
        except Exception as e:
            logger.error(f"Error processing case {case}: {str(e)}")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_filename = f"cwe_top25_analysis_filtered_results_{timestamp}.json"
    with open(result_filename, 'w') as f:
        json.dump(results, f, indent=4)

    logger.info(f"CWE Top 25 analysis process with AI filtering completed. Total files analyzed: {file_count}, Total CWE Top 25 issues found after filtering: {total_cwe_issues}")

if __name__ == "__main__":
    main()