From d1071d9d30c22a5a07bac11566d9c306b612a0b6 Mon Sep 17 00:00:00 2001
From: Sai Ashwin <saiashwin254@gmail.com>
Date: Thu, 14 Aug 2025 22:34:43 +0530
Subject: [PATCH] Update add_label_to_malware.py

---
 examples/add_label_to_malware.py | 77 +++++++++++++++++++++++---------
 1 file changed, 57 insertions(+), 20 deletions(-)

diff --git a/examples/add_label_to_malware.py b/examples/add_label_to_malware.py
index 48daa11a5..ef2903ef5 100644
--- a/examples/add_label_to_malware.py
+++ b/examples/add_label_to_malware.py
@@ -1,28 +1,65 @@
-# coding: utf-8
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+Added CLI argument parsing using argparse
+Added duplicate checks for malware and labels
+Added logging for better debugging
+Improved maintainability and usability of the script
 
+
+import argparse
+import logging
+import sys
 from pycti import OpenCTIApiClient
 
-# Variables
-api_url = "http://opencti:4000"
-api_token = "bfa014e0-e02e-4aa6-a42b-603b19dcf159"
 
-# OpenCTI initialization
-opencti_api_client = OpenCTIApiClient(api_url, api_token)
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+
+def main():
+    
+    parser = argparse.ArgumentParser(description="Add a label to a malware in OpenCTI.")
+    parser.add_argument("--url", required=True, help="OpenCTI API URL")
+    parser.add_argument("--token", required=True, help="OpenCTI API token")
+    parser.add_argument("--malware", required=True, help="Malware name")
+    parser.add_argument("--description", default="No description provided", help="Malware description")
+    parser.add_argument("--label", required=True, help="Label value to add")
+    parser.add_argument("--color", default="#ffa500", help="Label color (default: orange)")
+    args = parser.parse_args()
+
+    try:
+        client = OpenCTIApiClient(args.url, args.token)
+
+        
+        existing_malware = client.malware.read(filters=[{"key": "name", "values": [args.malware]}])
+        if existing_malware:
+            logging.info(f"Malware '{args.malware}' already exists.")
+            malware_id = existing_malware["id"]
+        else:
+            logging.info(f"Creating malware '{args.malware}'.")
+            malware = client.malware.create(name=args.malware, description=args.description)
+            malware_id = malware["id"]
+
+        
+        existing_label = client.label.read(filters=[{"key": "value", "values": [args.label]}])
+        if existing_label:
+            label_id = existing_label["id"]
+            logging.info(f"Label '{args.label}' already exists.")
+        else:
+            logging.info(f"Creating label '{args.label}'.")
+            label = client.label.create(value=args.label, color=args.color)
+            label_id = label["id"]
 
-# Create the malware
-malware = opencti_api_client.malware.create(
-    name="My new malware", description="A new evil tool."
-)
+        
+        logging.info(f"Adding label '{args.label}' to malware '{args.malware}'.")
+        client.stix_domain_object.add_label(id=malware_id, label_id=label_id)
 
-# Create the tag (if not exists)
-label = opencti_api_client.label.create(
-    value="Ransomware",
-    color="#ffa500",
-)
+        
+        malware = client.malware.read(id=malware_id)
+        logging.info("Updated Malware:")
+        print(malware)
 
-# Add the tag
-opencti_api_client.stix_domain_object.add_label(id=malware["id"], label_id=label["id"])
+    except Exception as e:
+        logging.error(f"Error: {e}")
+        sys.exit(1)
 
-# Print
-malware = opencti_api_client.malware.read(id=malware["id"])
-print(malware)
+if __name__ == "__main__":
+    main()