Merge pull request #26 from mayaCostantini/openai-wrapper

Add example OpenAI wrapper analytic for ranking suspicious processes

analytics/openai-suspicious-processes/Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.9
+
+RUN pip install --upgrade pip && \
+    pip install --no-cache-dir pandas openai
+
+WORKDIR /opt/analytics
+
+ADD analytics.py /opt/analytics/
+
+CMD ["python3", "analytics.py"]

analytics/openai-suspicious-processes/README.md

@@ -0,0 +1,20 @@
+# OpenAI suspicious processes ranking
+
+## Goal
+
+Analyze a list of potentially suspicious processes by prompting an OpenAI model.
+
+## Usage
+
+The example `PROMPT` in this analytic asks ChatGPT (`gpt-3.5-turbo` version) to rank the processes provided in the input dataframe and
+to give explanations for the top 10 suspicious processes.
+
+## Example
+
+Assume you have a variable `procs` with `process` entities. You can prompt the model to rank the processes with the following command:
+
+```
+APPLY python://openai-suspicious-processes ON procs
+```
+
+The `OPENAI_API_KEY` variable must be provided in the huntbook environment.

analytics/openai-suspicious-processes/analytics.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+import os
+import pandas as pd
+
+from openai import OpenAI
+
+
+OPENAI_MODEL = "gpt-3.5-turbo"
+
+PROMPT = """
+    The following dataframe contains information about different processes running on a system: {}. 
+    Rank those processes by suspicioussness and give an explanation for the top 10.
+    Focus on the `name` and `command_line` attributes of the processes.
+    """
+
+client = OpenAI(
+        api_key=os.environ.get("OPENAI_API_KEY"),
+    )
+
+
+def analytics(df):
+    """
+    Given a prompt and process information in the dataframe,
+    rank the processes by suspiciousness and give an explanation
+    for the top 10 suspicious processes.
+    """
+
+    complete_prompt = PROMPT.format(df.to_json())
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": complete_prompt,
+            }
+        ],
+        model=OPENAI_MODEL,
+    )
+
+    display = (
+                f"<p><b>Prompt:</b> {PROMPT.format(df)} </p>"
+                f"<p><b>Answer:</b> {chat_completion.choices[0].message.content}</p>"
+              )
+
+    return df, display