Introducing new team member Quark Script Agent (#37)

Author: pulorsok

README.md

@@ -1,33 +1,53 @@
-# Quickstart Quark Script
+# Quark Script Agent
 
-In this tutorial, we will learn how to install and run Quark Script with a very easy example.
-We show how to detect CWE-798 in ovaa.apk. 
+Introducing Quark's new member, the Quark Script Agent, the first AI assistant in the Quark team. This agent enables users to perform analyses using natural language, without the need for programming or scripting expertise, making the process simple and user-friendly.
 
-### STEP 1: Environments Requirements
-* Quark Script requires Python 3.8+
+The Quark Script Agent integrates with LangChain, which utilizes OpenAI's large language models to act as a bridge between natural language and the Quark Script API. LangChain defines the Quark Script API as a tool that large language models can understand and use. This means that users can easily call new analysis APIs using natural language commands by simply adding new tools as needed.
 
-### STEP 2: Install Quark Engine
-You can install Quark Engine by running:
-  ```
-  pip3 install quark-engine
-  ```
+## Showcase: Detecting CWE-798 with Quark Script Agent
+Here's an example of using the Quark Script Agent with the `quarkScriptAgent.py`. This agent can currently detect [CWE-798](https://cwe.mitre.org/data/definitions/798.html) vulnerability in the [ovaa.apk](https://github.com/oversecured/ovaa). See the details below.
 
-### STEP 3: Prepare Quark Script, Detection Rule and the Sample File
-1. Get the CWE-798 Quark Script and the detection rule [here](https://quark-engine.readthedocs.io/en/latest/quark_script.html#detect-cwe-798-in-android-application-ovaa-apk).
-2. Get the sampe file (ovaa.apk) [here](https://github.com/dark-warlord14/ovaa/releases/tag/1.0).
-3. Put the script, detection rule, and sample file in the same directory.
-4. Edit accordingly to the file names:
-```python
-SAMPLE_PATH = "ovaa.apk"
-RULE_PATH = "findSecretKeySpec.json"
+### Quick Start
+
+1. clone the repository:
+```
+git clone https://github.com/quark-engine/quark-script.git
+```
+
+2. Install the required packages:
+```
+pip install -r requirements.txt
 ```
 
-### STEP 4: Run the script
+3. Run the script:
 ```
-python3 CWE-798.py
+python quarkScriptAgent.py
 ```
 
-You should now see the detection result in the terminal:
+4. Result:
+
+<img width="1440" alt="截圖 2024-07-26 下午3 39 12" src="https://github.com/user-attachments/assets/9c8ba9d3-c8b5-4583-8cb8-750f8c3bf2a7">
+
+### Decode the Prompts
+Here are two prompts, each for executing different analysis processes.
+
 ```
-Found hard-coded AES key 49u5gh249gh24985ghf429gh4ch8f23f
+1st Prompt: Initialize the rule instance with the rule path set to "rule.json"
 ```
+Used Quark Script APIs/Tools that LLM used:  `Rule()`
+
+```
+2nd Prompt: Run Quark Analysis using the rule instance on the apk sample "ovaa.apk", 
+            and Check if the parameters are hard-coded. If yes, display the hard-coded values.
+```
+Used Quark Script APIs/Tools that LLM used: `runQuarkAnalysis()`, `getParameterValues()` and `isHardCoded()`
+
+The `Rule()`, `runQuarkAnalysis()`, `getParameterValues()`, and `isHardCoded()` functions are treated as **tools** within LangChain, enabling them to be invoked through the `gpt-4o` model to analyze and identify [CWE-798](https://cwe.mitre.org/data/definitions/798.html) vulnerabilities in the [ovaa.apk](https://github.com/oversecured/ovaa) sample.
+
+<img width="829" alt="截圖 2024-07-26 下午9 25 23" src="https://github.com/user-attachments/assets/14de8563-e52e-4bdc-9960-ec73cbd10ada">
+
+
+* Notes: 
+  1. Since LangChain currently does not support passing Python instances between tools, we are temporarily using global variables to pass parameters between tools in `quarkScriptAgent.py`.
+  2. Place the rules, samples, and `quarkScriptAgent.py` in the same folder; the LLM will automatically find files with matching names.
+  3. A web GUI is under construction, please stay tuned!

ovaa.apk

@@ -0,0 +1,144 @@
+import os
+import re
+
+from termcolor import colored
+
+from langchain_openai import ChatOpenAI
+from langchain.agents import tool, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
+from langchain.agents.format_scratchpad.openai_tools import (
+    format_to_openai_tool_messages,
+)
+
+from quark.script import Rule, _getQuark, QuarkResult
+
+if "OPENAI_API_KEY" not in os.environ:
+    api_key = input("OpenAI API Key: ")
+    os.environ["OPENAI_API_KEY"] = api_key
+
+
+@tool
+def loadRule(rulePath: str):
+    """
+    Given a rule path,
+    this instance loads a rule from the rule path.
+    """
+
+    global ruleInstance
+    ruleInstance = Rule(rulePath)
+
+    return "Rule defined successfully"
+
+
+@tool
+def runQuarkAnalysis(samplePath: str):
+    """
+    Given detection rule and target sample,
+    this instance runs the Quark Analysis.
+    """
+
+    global ruleInstance
+    global quarkResultInstance
+
+    quark = _getQuark(samplePath)
+    quarkResultInstance = QuarkResult(quark, ruleInstance)
+
+    return "Quark analysis completed successfully"
+
+
+@tool
+def getBehaviorOccurList():
+    """
+    Given the Quark analysis result,
+    this instance extracts the behavior occurrence list.
+    """
+
+    global quarkResultInstance
+    global behaviorOccurList
+
+    behaviorOccurList = quarkResultInstance.behaviorOccurList
+    return "Behavior occurrence list extracted successfully"
+
+
+@tool
+def getParameterValues():
+    """
+    Given the behavior occurrence list,
+    this instance extracts the parameter values.
+    """
+
+    global behaviorOccurList
+    global parameters
+
+    for behavior in behaviorOccurList:
+        param = behavior.getParamValues()[1]
+
+    parameters = re.findall(r"\((.*?)\)", param)[1]
+
+    return "Parameter values extracted successfully"
+
+
+@tool
+def isHardCoded():
+    """
+    Given the parameter values,
+    this instance checks if the parameter values are hard-coded.
+    """
+
+    global parameters
+    global quarkResultInstance
+
+    # check parameter values are hard-coded
+    if parameters in quarkResultInstance.getAllStrings():
+        return parameters
+
+    return False
+
+
+tools = [
+    loadRule,
+    runQuarkAnalysis,
+    getBehaviorOccurList,
+    getParameterValues,
+    isHardCoded,
+]
+
+
+llm = ChatOpenAI(model="gpt-4o", temperature=0.1)
+llm_with_tools = llm.bind_tools(tools)
+
+prompt = ChatPromptTemplate.from_messages([
+    (
+        "system",
+        "You are very powerful assistant, but don't know current events",
+    ),
+    ("user", "{input}"),
+    MessagesPlaceholder(variable_name="agent_scratchpad"),
+])
+
+agent = (
+    {
+        "input": lambda x: x["input"],
+        "agent_scratchpad": lambda x: format_to_openai_tool_messages(
+            x["intermediate_steps"]
+        ),
+    }
+    | prompt
+    | llm_with_tools
+    | OpenAIToolsAgentOutputParser()
+)
+
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
+
+input_text = input(colored('User Input: ', 'green'))
+while input_text.lower() != 'bye':
+    if input_text:
+        response = agent_executor.invoke({
+            'input': input_text,
+        })
+        print()
+        print(colored('Agent: ', "cyan"), response['output'])
+        print()
+
+        input_text = input(colored('User Input: ', 'green'))

quarkScriptAgent.py

@@ -0,0 +1,5 @@
+termcolor==2.4.0
+langchain==0.2.11
+langchain-core==0.2.23
+langchain-openai==0.1.17
+quark-engine==24.7.1

requirement.txt

@@ -0,0 +1,18 @@
+{
+    "crime": "Detect APK using SecretKeySpec.",
+    "permission": [],
+    "api": [
+        {
+            "descriptor": "()[B",
+            "class": "Ljava/lang/String;",
+            "method": "getBytes"
+        },
+        {
+            "descriptor": "([BLjava/lang/String;)V",
+            "class": "Ljavax/crypto/spec/SecretKeySpec;",
+            "method": "<init>"
+        }
+    ],
+    "score": 1,
+    "label": []
+}