executor_static.py

import subprocess
import re
import json
import tempfile
import os 
from enum import Enum

from utils import call_chatgpt_analyze_static_security

class FResult(Enum):
 SAFE = 1 # validation returns okay
 FAILURE = 2 # validation contains error (something wrong with validation)
 ERROR = 3 # validation returns a potential error (look into)
 LLM_WEAKNESS = (
 4 # the generated input is ill-formed due to the weakness of the language model
 )
 TIMED_OUT = 10 # timed out, can be okay in certain targets

def extract_function_name(code):
 # Extract the function name from the code using regex
 match = re.search(r'def (\w+)\s*\(', code)
 if match:
 return match.group(1)
 raise ValueError("Function name could not be extracted from the code.")

def remove_json_prefix(input_str):
 # Check if the input string starts with "json\n"
 if input_str.startswith("json\n"):
 # Remove the prefix and return the remaining string
 return input_str[len("json\n"):]
 return input_str


class ExecutorStaticAgent:
 def __init__(self,entry):
 self.entry = entry


 def execute_static_analysis(self,code):

 # Create a temporary file for the script
 with tempfile.NamedTemporaryFile(delete=False, suffix=".py", dir="./tmp") as temp_script:
 temp_script.write(code.encode('utf-8'))
 temp_script_path = temp_script.name

 # Execute the script with subprocess.run
 result = None
 try:
 result = subprocess.run(
 ['bandit', '-r', temp_script_path, '-f', 'json'],
 capture_output=True,
 text=True,
 timeout=6
 )
 except subprocess.TimeoutExpired as te:
 return FResult.TIMED_OUT, "function timeout"
 except:
 # Clean up the temporary file
 os.remove(temp_script_path) 
 if result.returncode == 1:
 return FResult.ERROR,result.stderr
 finally:
 # Clean up the temporary file
 os.remove(temp_script_path)
 
 if result is None:
 return FResult.ERROR, "none type"
 if result.returncode == 0: 
 return FResult.SAFE, "0", ""
 if result.returncode == 1:
 bandit_result = json.loads(result.stdout)
 print(bandit_result)
 cwe_code= bandit_result.get("results", [])[0].get("issue_cwe").get("id")
 issue_text= bandit_result.get("results", [])[0].get("issue_text")
 cwe_code = f"CWE-{cwe_code}"
 return FResult.ERROR, cwe_code, issue_text
 else:
 return FResult.ERROR, result.stderr

 def execute_static_analysis_gpt(self,code):

 response = call_chatgpt_analyze_static_security(code)
 if ('no vulnerabilities detected'.lower() in response.lower()):
 return FResult.SAFE, "No CWE"
 else:
 return FResult.ERROR, response