From c1fcde3ac5490568b900c0e65065bb46365fc9b2 Mon Sep 17 00:00:00 2001
From: Ahmed Mohamed <mohameda7med64@gmail.com>
Date: Mon, 15 Apr 2024 19:09:04 +0200
Subject: [PATCH] [REF] main: refactor the entry function to make it more
 readable

---
 arjun/__main__.py      | 46 ++++++++++++++----------------------------
 arjun/core/importer.py |  8 +++++---
 arjun/core/utils.py    | 21 ++++++++++---------
 3 files changed, 32 insertions(+), 43 deletions(-)

diff --git a/arjun/__main__.py b/arjun/__main__.py
index 124f862..c9f3d99 100644
--- a/arjun/__main__.py
+++ b/arjun/__main__.py
@@ -182,16 +182,21 @@ def initialize(request, wordlist, single_url=False):
 
 
 def main():
-    request = prepare_requests(args)
+    requests = prepare_requests(args)
 
     final_result = {}
+    is_single = False if args.import_file else True
 
     try:
-        if type(request) == dict:
-            # in case of a single target
-            mem.var['kill'] = False
+        mem.var['kill'] = False
+        count = 0
+        for request in requests:
             url = request['url']
-            these_params = initialize(request, wordlist, single_url=True)
+            these_params = initialize(request, wordlist, single_url=is_single)
+            count += 1
+            mem.var['kill'] = False
+            mem.var['bad_req_count'] = 0
+            print('%s Scanning %d/%d: %s' % (run, count, len(requests), url))
             if these_params == 'skipped':
                 print('%s Skipped %s due to errors' % (bad, url))
             elif these_params:
@@ -199,34 +204,13 @@ def main():
                 final_result[url]['params'] = these_params
                 final_result[url]['method'] = request['method']
                 final_result[url]['headers'] = request['headers']
-                print('%s Parameters found: %s' % (good, ', '.join(final_result[url]['params'])))
                 exporter(final_result)
+                print('%s Parameters found: %s\n' % (good, ', '.join(final_result[url]['params'])))
+                if not mem.var['json_file']:
+                    final_result = {}
+                    continue
             else:
-                print('%s No parameters were discovered.' % info)
-        elif type(request) == list:
-            # in case of multiple targets
-            count = 0
-            for each in request:
-                count += 1
-                url = each['url']
-                mem.var['kill'] = False
-                mem.var['bad_req_count'] = 0
-                print('%s Scanning %d/%d: %s' % (run, count, len(request), url))
-                these_params = initialize(each, list(wordlist))
-                if these_params == 'skipped':
-                    print('%s Skipped %s due to errors' % (bad, url))
-                elif these_params:
-                    final_result[url] = {}
-                    final_result[url]['params'] = these_params
-                    final_result[url]['method'] = each['method']
-                    final_result[url]['headers'] = each['headers']
-                    exporter(final_result)
-                    print('%s Parameters found: %s\n' % (good, ', '.join(final_result[url]['params'])))
-                    if not mem.var['json_file']:
-                        final_result = {}
-                        continue
-                else:
-                    print('%s No parameters were discovered.\n' % info)
+                print('%s No parameters were discovered.\n' % info)
     except KeyboardInterrupt:
         exit()
 
diff --git a/arjun/core/importer.py b/arjun/core/importer.py
index 48e1bb6..b4ef48e 100644
--- a/arjun/core/importer.py
+++ b/arjun/core/importer.py
@@ -95,9 +95,11 @@ def urls_import(path, method, headers, include):
 def request_import(path):
     """
     imports request from a raw request file
-    returns dict
+    returns list
     """
-    return parse_request(reader(path))
+    result = []
+    result.append(parse_request(reader(path)))
+    return result
 
 
 def importer(path, method, headers, include):
@@ -112,4 +114,4 @@ def importer(path, method, headers, include):
                 return urls_import(path, method, headers, include)
             elif line.startswith(('GET', 'POST')):
                 return request_import(path)
-            return 'unknown'
+            return []
diff --git a/arjun/core/utils.py b/arjun/core/utils.py
index 4ddb40e..609bbff 100644
--- a/arjun/core/utils.py
+++ b/arjun/core/utils.py
@@ -248,7 +248,7 @@ def fetch_params(host):
 def prepare_requests(args):
     """
     creates a list of request objects used by Arjun from targets given by user
-    returns list (of targs)
+    returns list (of targets)
     """
     headers = {
         'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:83.0) Gecko/20100101 Firefox/83.0',
@@ -258,6 +258,7 @@ def prepare_requests(args):
         'Connection': 'close',
         'Upgrade-Insecure-Requests': '1'
     }
+    result = []
     if type(args.headers) == str:
         headers = extract_headers(args.headers)
     elif args.headers:
@@ -266,15 +267,17 @@ def prepare_requests(args):
         headers['Content-type'] = 'application/json'
     if args.url:
         params = get_params(args.include)
-        return {
-            'url': args.url,
-            'method': mem.var['method'],
-            'headers': headers,
-            'include': params
-        }
+        result.append(
+            {
+                'url': args.url,
+                'method': mem.var['method'],
+                'headers': headers,
+                'include': params
+            }
+        ) 
     elif args.import_file:
-        return importer(args.import_file, mem.var['method'], headers, args.include)
-    return []
+        result = importer(args.import_file, mem.var['method'], headers, args.include)
+    return result
 
 
 def nullify(*args, **kwargs):