From 953f609f10a099b084e09baeb6c1061172d66395 Mon Sep 17 00:00:00 2001
From: Sanket Baraiya <56958135+SanketBaraiya@users.noreply.github.com>
Date: Tue, 7 Oct 2025 19:03:10 +0530
Subject: [PATCH] refactor: use env for API key and validate VirusTotal API key

---
 amcache-evilhunter.py | 33 +++++++++++++++++++++++++++------
 env                   |  2 ++
 requirements.txt      |  1 +
 3 files changed, 30 insertions(+), 6 deletions(-)
 create mode 100644 env

diff --git a/amcache-evilhunter.py b/amcache-evilhunter.py
index 629705f..fc54d91 100644
--- a/amcache-evilhunter.py
+++ b/amcache-evilhunter.py
@@ -14,6 +14,7 @@
 from pathlib import Path
 from functools import lru_cache
 from datetime import datetime, timedelta
+from dotenv import dotenv_values
 
 import requests
 from requests.exceptions import HTTPError
@@ -115,6 +116,7 @@ def missing_publisher(data):
 
 class AmcacheParser:
     """Parser for offline Amcache.hve registry hive."""
+
     def __init__(self, hive_path, start=None, end=None):
         if not hive_path.exists():
             raise FileNotFoundError(f"Hive file not found: {hive_path}")
@@ -395,9 +397,9 @@ def main():
 
     group = parser.add_mutually_exclusive_group()
     group.add_argument("-v", "--vt", action="store_true",
-                       help="Enable VirusTotal lookups (requires VT_API_KEY)")
+                       help="Enable VirusTotal lookups (requires VT_API_KEY in .env)")
     group.add_argument("--opentip", action="store_true",
-                       help="Enable Kaspersky OpenTIP lookups (requires OPENTIP_API_KEY)")
+                       help="Enable Kaspersky OpenTIP lookups (requires OPENTIP_API_KEY in .env)")
 
     parser.add_argument('-V', '--version', action='version',
                         version=f"AmCache-EvilHunter {VERSION} by Cristian Souza")
@@ -416,17 +418,36 @@ def main():
     parser.add_argument("--csv", type=Path, help="Path to write CSV")
     args = parser.parse_args()
 
+    env_vars = dotenv_values("env")
+
     vt_api_key = None
     ot_api_key = None
     if args.vt:
-        vt_api_key = os.getenv("VT_API_KEY")
+        vt_api_key = env_vars.get("VT_API_KEY")
         if not vt_api_key:
-            console.print("[bold red]Error:[/] VT_API_KEY environment variable not set", style="red")
+            console.print("[bold red]Error:[/] VT_API_KEY not set in .env", style="red")
             sys.exit(1)
+        else:
+            url = "https://www.virustotal.com/api/v3/users/me"
+            headers = {"x-apikey": vt_api_key}
+
+            response = requests.get(url, headers=headers)
+
+            if response.status_code == 200:
+                console.print("[bold green]Success:[/] VirusTotal API key is valid!", style="green")
+                user_info = response.json()
+            elif response.status_code == 401:
+                console.print(f"[bold red]Error:[/] Invalid VirusTotal API key. [bold red]Reason:[/] {response.json().get('error').get('message')}", style="red")
+                sys.exit(1)
+            else:
+                console.print(f"[bold red]Error:[/] Checking API key: {response.status_code}", style="red")
+                console.print(response.text)
+                sys.exit(1)
+
     if args.opentip:
-        ot_api_key = os.getenv("OPENTIP_API_KEY")
+        ot_api_key = env_vars.get("OPENTIP_API_KEY")
         if not ot_api_key:
-            console.print("[bold red]Error:[/] OPENTIP_API_KEY environment variable not set", style="red")
+            console.print("[bold red]Error:[/] OPENTIP_API_KEY not set in .env", style="red")
             sys.exit(1)
 
     # parse date filters
diff --git a/env b/env
new file mode 100644
index 0000000..c4ffedb
--- /dev/null
+++ b/env
@@ -0,0 +1,2 @@
+VT_API_KEY=YOUR_API_KEY
+OPENTIP_API_KEY=YOUR_API_KEY
diff --git a/requirements.txt b/requirements.txt
index 5fbf6cc..baed7b9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 requests==2.32.4
 python-registry==1.3.1
 rich==14.0.0
+python-dotenv