cwe_top25_filter2.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_improved_{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. Consider the full context of the code, including potential security measures already in place.
3. Be aware that using environment variables for storing sensitive information is generally considered secure.
4. Only report violations if you are highly confident they exist and are not false positives.
5. If you find any violations, provide the following information for each:
   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, including the specific lines of code involved
   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, including line numbers",
            "mitigation": "Suggested mitigation strategy"
        }}
    ]
}}

If no violations are found or you're not highly confident about any potential violations, return an empty "risks" array.
"""
    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. Consider the full context and any security measures already in place.
3. Be aware that using environment variables for storing sensitive information is generally considered secure.
4. Only include a risk if you are highly confident it is a true positive.
5. For each true positive, provide a detailed explanation of why it's a genuine security risk, including specific code lines.

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

If all risks are determined to be false positives or you're not highly confident about any of them, return an empty list for "filtered_risks".
"""
    messages = [{"role": "user", "content": prompt}]
    return call_llm(messages)

def additional_checks(code_path, filtered_risks):
    # Implement additional checks here to further reduce false positives
    # For example, you could add specific checks for certain CWE types
    return filtered_risks

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, total_filtered_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)}")

    initial_risks = parsed_result['risks']
    initial_risk_count = len(initial_risks)

    if not initial_risks:
        logger.info("No risks found in initial analysis. Skipping false positive filtering.")
        final_risks = []
        filtered_count = 0
    else:
        logger.info("Filtering false positives...")
        filtered_result = filter_false_positives(before_code_path, initial_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)}")

        final_risks = additional_checks(before_code_path, filtered_parsed_result['filtered_risks'])
        filtered_count = initial_risk_count - len(final_risks)

    file_count += 1
    cwe_issues = len(final_risks)
    total_cwe_issues += cwe_issues
    total_filtered_issues += filtered_count

    logger.info(f"Files analyzed: {file_count}")
    logger.info(f"Initial CWE Top 25 issues found: {initial_risk_count}")
    logger.info(f"Issues filtered out: {filtered_count}")
    logger.info(f"Final CWE Top 25 issues: {cwe_issues}")
    logger.info(f"Total issues found so far: {total_cwe_issues}")
    logger.info(f"Total issues filtered out so far: {total_filtered_issues}")

    return {"filtered_risks": final_risks}, file_count, total_cwe_issues, total_filtered_issues

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

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

    results = {}
    file_count = 0
    total_cwe_issues = 0
    total_filtered_issues = 0
    
    for case in cases:
        try:
            case_result, file_count, total_cwe_issues, total_filtered_issues = process_case(case, logger, file_count, total_cwe_issues, total_filtered_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 all checks")
            else:
                logger.info(f"Case {case}: {risks_count} CWE Top 25 risks identified after all checks")
            
        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_improved_results_{timestamp}.json"
    with open(result_filename, 'w') as f:
        json.dump(results, f, indent=4)

    logger.info(f"Improved CWE Top 25 analysis process completed.")
    logger.info(f"Total files analyzed: {file_count}")
    logger.info(f"Total CWE Top 25 issues found after all checks: {total_cwe_issues}")
    logger.info(f"Total issues filtered out: {total_filtered_issues}")
    
    if total_filtered_issues > 0:
        filter_effectiveness = (total_filtered_issues / (total_cwe_issues + total_filtered_issues)) * 100
        logger.info(f"Filter effectiveness: {filter_effectiveness:.2f}% of initially identified issues were filtered out")
    else:
        logger.info("No issues were filtered out. Consider if the filtering step is necessary.")

if __name__ == "__main__":
    main()