Introducing the Fortify rule library

audit/__init__.py

@@ -1,8 +1,6 @@
 import os
 import re
-import time
 import uuid
-import tiktoken
 import xml.etree.ElementTree as ET
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_core.messages import SystemMessage
@@ -20,17 +18,18 @@ from audit import callback
 from audit.prompt import SYSTEM_PROMPT
 from audit.language import LANGUAGE
 
-reasoning_model = 'gemini-2.0-flash-thinking-exp'
+reasoning_model = 'gpt-4o'
 embedding_model = 'text-embedding-3-large'
 
 xml_pattern = r'<root>.*?</root>'
 
 
 class Audit:
-    def __init__(self):
+    def __init__(self, fortify_rules):
         self.raw_chain = None
         self.source_files_list = []
         self.max_token = 4096
+        self.fortify_rules = fortify_rules
         self.chat_history = ChatMessageHistory()
         self.session_id = uuid.uuid4().hex
         self.response_callback = callback.CustomCallbackHandler()
@@ -93,12 +92,18 @@ class Audit:
                     self.log.info(f'Request source code: {content}')
                     input_content = open(content, 'r', encoding='utf-8').read()
                     continue
+                elif action == 'QUERY FORTIFY':
+                    self.log.info(f'Request fortify: {content}')
+                    input_content = '\n'.join(x for x in self.fortify_rules if x == content)
+                    continue
                 elif action == 'OUTPUT RESULT':
                     self.log.warning(f'Audit result: \n\n{content}')
                     self.store_messages_in_faiss(content)
                     callback_function(content)  # Callback function, used to obtain results externally
                     input_content = ''
                     continue
+                elif action == 'FINISH TASK':
+                    self.log.info(content)
                 else:
                     self.log.critical(f'Unknown action! {action}')
                     break
@@ -126,8 +131,6 @@ class Audit:
             history_messages_key='messages',
         )
 
-        self.log.debug(f'Chat messages: {input_dict}')
-
         for _ in chain_with_message_history.stream(input_dict, config_dict):
             pass
 

audit/language.py

@@ -1,8 +1,9 @@
 LANGUAGE = {
     'c': ['.c'],
-    'c++': ['.cc', '.cpp'],
+    'cpp': ['.cc', '.cpp'],
     'go': ['.go'],
     'php': ['php', 'php3', 'php4', 'php5', 'phtml'],
+    'jsp': ['.jsp'],
     'java': ['.java'],
     'python': ['.py'],
     'javascript': ['.js'],

audit/prompt.py

@@ -1,7 +1,15 @@
 SYSTEM_PROMPT = """
 You are an intelligent code auditor. I will provide you with a source code. Please strictly follow the following requirements to conduct code audit.
-All your output must strictly follow the following specifications. It is forbidden to output in any other form (including plain text, Markdown, etc.). 
-and it is forbidden to bring "`" when outputting:
+During the audit process, you can refer to Fortify's rule base(Execute Action 3), but it does not have to be completely consistent to determine the existence of a vulnerability. The rule base format provided to you is as follows:
+{
+    'language':
+    'vuln_kingdom':
+    'vuln_category':
+}
+
+Before officially starting the audit, it is recommended to query the Fortify rule base as a reference.
+All your output must strictly follow the following specifications. It is forbidden to output in any other form (including plain text, Markdown, etc.), and it is forbidden to bring "`" when outputting.
+You can choose to perform the following actions:
 
 1. Query project structure:
 <root>
@@ -15,13 +23,27 @@ and it is forbidden to bring "`" when outputting:
 <content>the absolute path of the file you want to query</content>
 </root>
 
-3. Output audit results
+3. Query fortify
+<root>
+<action>QUERY FORTIFY</action>
+<content>The language you want to query, options are: c, cpp, go, php, jsp, java, python, javascript</content>
+</root>
+
+4. Output audit results
 <root>
 <action>OUTPUT RESULT</action>
 <content>the audit results you want to output</content>
 </root>
 
-4. End the audit task
+The output result format is as follows(JSON):
+{
+    "Vulnerability Type": 
+    "Vulnerability File": 
+    "Vulnerability Code Summary": 
+    "Vulnerability repair suggestions":
+} 
+
+5. End the audit task
 <root>
 <action>FINISH TASK</action>
 <content></content>