Merge pull request #14 from ackatz/v0.5.0

V0.5.0

.github/workflows/ci.yml

@@ -37,6 +37,9 @@ jobs:
         
         echo "[threatfox]" >> config.ini
         echo "api_key = ${{ secrets.threatfox_api_key }}" >> config.ini
+        
+        echo "[pulsedive]" >> config.ini
+        echo "api_key = ${{ secrets.pulsedive_api_key }}" >> config.ini
 
     - name: Install dependencies
       run: |

README.md

@@ -30,6 +30,12 @@ For example, to look up IP 1.1.1.1 on AbuseIPDB, you'd run:
 seclook abuseipdb 1.1.1.1
 ```
 
+You can send the JSON response to OpenAI GPT-4 for summarization:
+
+```bash
+seclook virustotal 44d88612fea8a8f36de82e1278abb02f --gpt4
+```
+
 You can pipe the output to `fx` or `jq` for further processing:
 
 ```bash
@@ -40,11 +46,12 @@ You can `grep` the output for known keys to get specific information:
 
 ```bash
 seclook virustotal 44d88612fea8a8f36de82e1278abb02f | grep malicious
-``` 
+```
 
 ## Options
 
-`--export` – Use this flag to export the results to a JSON file on your Desktop.
+`--export` – Use this flag to export the results to a JSON file on your Desktop.  
+`--gpt4` – Use this flag to summarize the JSON response from a service in GPT4.
 
 ## Supported services
 

config.ini.sample

@@ -21,3 +21,6 @@ api_key =
 
 [virustotal]
 api_key =
+
+[openai]
+api_key =

seclook/cli.py

@@ -10,6 +10,7 @@
     threatfox_lookup,
     pulsedive_lookup,
 )
+from seclook.openai import gpt4_summarize
 import json
 import os
 
@@ -18,7 +19,8 @@
 @click.argument("service", required=False)
 @click.argument("value", required=False)
 @click.option("--export", is_flag=True, help="Export JSON to Desktop")
-def main(service, value, export):
+@click.option("--gpt4", is_flag=True, help="Use GPT-4 to summarize results")
+def main(service, value, export, gpt4):
     """Perform lookups from various security services
 
     - Use `seclook [service] [value]` to perform a lookup.
@@ -81,6 +83,11 @@ def main(service, value, export):
         click.echo(f"Results exported to {filename}")
         return
 
+    if gpt4:
+        result = gpt4_summarize.search(service, result)
+        pretty_result = json.dumps(result, indent=4)
+        click.echo(pretty_result)
+
     pretty_result = json.dumps(result, indent=4)
     click.echo(pretty_result)
 

seclook/openai/gpt4_summarize.py

@@ -0,0 +1,26 @@
+import requests
+from seclook.config_helper import load_config
+
+base_url = "https://api.openai.com/v1/chat/completions"
+
+
+def search(service, json_response):
+    config, openai_api_key = load_config("openai")
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer " + openai_api_key,
+    }
+    payload = {
+        "model": "gpt-4",
+        "messages": [
+            {
+                "role": "user",
+                "content": f"You are a security analyst summarizer. "
+                f"Given the following JSON response from {service}"
+                f", summarize the response to tell me what I need to know: "
+                f"{json_response}",
+            }
+        ],
+    }
+    response = requests.post(base_url, headers=headers, json=payload)
+    return response.json()

setup.py

@@ -5,7 +5,7 @@
     author="Andrew Katz",
     author_email="[email protected]",
     license="MIT",
-    version="0.4.0",
+    version="0.5.0",
     packages=find_packages(),
     include_package_data=True,
     install_requires=[