From 8ece46e8025b0715a0197f4e9ec7d3b76635cf23 Mon Sep 17 00:00:00 2001
From: 0n1cOn3 <0n1cOn3@gmx.ch>
Date: Sat, 21 Oct 2023 16:17:13 +0200
Subject: [PATCH] refactored code

---
 search.py | 61 +++++++++++++++++++++++++++++++------------------------
 1 file changed, 34 insertions(+), 27 deletions(-)

diff --git a/search.py b/search.py
index 5304fea..92417a4 100644
--- a/search.py
+++ b/search.py
@@ -4,47 +4,54 @@
 import argparse
 from dotenv import load_dotenv
 
-load_dotenv()
-
-SHODAN_API_KEY = os.environ.get("API_KEY")
-
-if SHODAN_API_KEY is None:
-    print("API key not found in environment variables.")
-    sys.exit(1)
-
-parser = argparse.ArgumentParser()
-subparsers = parser.add_subparsers(dest="command", required=True)
-
-search_parser = subparsers.add_parser("search")
-search_parser.add_argument("query", type=str)
-
-scan_parser = subparsers.add_parser("scan")
-scan_parser.add_argument("file", type=str)
-
-args = parser.parse_args()
-
-if args.command == "search":
-    api = shodan.Shodan(SHODAN_API_KEY)
+def search(query):
     try:
-        results = api.search(args.query)
+        results = api.search(query)
         with open("ip-temp.txt", "w") as f:
             for result in results["matches"]:
                 ip = result["ip_str"]
                 print("Found IP:", ip)
                 f.write(ip + "\n")
     except shodan.exception.APIError as e:
-        print("Error occurred:", e)
+        print("Error occurred during search:", e)
         sys.exit(1)
-elif args.command == "scan":
+
+def scan(file):
     try:
-        with open(args.file, "r") as f, open("results.txt", "w") as results_file:
+        with open(file, "r") as f, open("results.txt", "w") as results_file:
             for i, line in enumerate(f):
                 ip = line.strip()
                 print("Scanning {} out of {} IPs: {}".format(i+1, len(f), ip))
                 os.system("nmap -sS {} >> results.txt".format(ip))
     except FileNotFoundError:
-        print("File not found:", args.file)
+        print("File not found:", file)
         sys.exit(1)
     except Exception as e:
-        print("Error occurred:", e)
+        print("Error occurred during scanning:", e)
         sys.exit(1)
+
+if __name__ == "__main__":
+    load_dotenv()
+    SHODAN_API_KEY = os.environ.get("API_KEY")
+
+    if SHODAN_API_KEY is None:
+        print("API key not found in environment variables.")
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    search_parser = subparsers.add_parser("search")
+    search_parser.add_argument("query", type=str)
+
+    scan_parser = subparsers.add_parser("scan")
+    scan_parser.add_argument("file", type=str)
+
+    args = parser.parse_args()
+
+    api = shodan.Shodan(SHODAN_API_KEY)
+
+    if args.command == "search":
+        search(args.query)
+    elif args.command == "scan":
+        scan(args.file)