cleanup; added flag to constructors

gato/attack/attack.py

@@ -38,15 +38,15 @@ def __init__(
         author_name: str = None,
         timeout: int = 30,
         github_url: str = None,
-	no_sleep: bool = False
+        no_sleep: bool = False
     ):
 
         self.api = Api(
             pat,
             socks_proxy=socks_proxy,
             http_proxy=http_proxy,
             github_url=github_url,
-	    no_sleep=no_sleep,
+            no_sleep=no_sleep,
         )
 
         self.socks_proxy = socks_proxy

gato/cli/cli.py

@@ -49,12 +49,6 @@ def cli(args):
         required=False,
     )
 
-    parser.add_argument(
-        "--no-sleep",
-        help="Disable sleeping when api rate limit is hit and exit instead",
-        action="store_true"
-    )
-
     attack_parser = subparsers.add_parser(
         "attack", help="CI/CD Attack Capabilities", aliases=["a"],
         formatter_class=argparse.RawTextHelpFormatter
@@ -170,7 +164,8 @@ def attack(args, parser):
         socks_proxy=args.socks_proxy,
         http_proxy=args.http_proxy,
         timeout=timeout,
-        github_url=args.api_url
+        github_url=args.api_url,
+        no_sleep=args.no_sleep
     )
 
     if args.pull_request:
@@ -233,7 +228,8 @@ def enumerate(args, parser):
             http_proxy=args.http_proxy,
             output_yaml=args.output_yaml,
             skip_log=args.skip_runlog,
-            github_url=args.api_url
+            github_url=args.api_url,
+            no_sleep=args.no_sleep
         )
 
     exec_wrapper = Execution()
@@ -351,6 +347,12 @@ def configure_parser_general(parser):
         action="store_true"
     )
 
+    parser.add_argument(
+        "--no-sleep",
+        help="Exit immediately upon the API Rate Limit being hit.",
+        action="store_true"
+    )
+
 
 def configure_parser_attack(parser):
     """Helper method to add arguments to the attack subparser.

gato/enumerate/enumerate.py

@@ -25,7 +25,7 @@ def __init__(
         skip_log: bool = False,
         github_url: str = None,
         output_json: str = None,
-	no_sleep: bool = False
+        no_sleep: bool = False
     ):
         """Initialize enumeration class with arguments sent by user.
 
@@ -47,7 +47,7 @@ def __init__(
             socks_proxy=socks_proxy,
             http_proxy=http_proxy,
             github_url=github_url,
-	    no_sleep=no_sleep,
+            no_sleep=no_sleep,
         )
 
         self.socks_proxy = socks_proxy

gato/github/api.py

@@ -6,6 +6,7 @@
 import zipfile
 import re
 import io
+import sys
 
 from gato.cli import Output
 from datetime import datetime, timezone, timedelta
@@ -28,7 +29,8 @@ class Api():
 
     def __init__(self, pat: str, version: str = "2022-11-28",
                  http_proxy: str = None, socks_proxy: str = None,
-                 github_url: str = "https://api.github.com", no_sleep: bool = False):
+                 github_url: str = "https://api.github.com",
+                 no_sleep: bool = False):
         """Initialize the API abstraction layer to interact with the GitHub
         REST API.
 
@@ -101,14 +103,14 @@ def __check_rate_limit(self, headers):
             # all calling code. We inform the here user that we are sleeping.
             # very large orgs will take several hours to enumerate, especially
             # if runlog enumeration is enabled.
+            if self.no_sleep:
+                Output.warn("Exiting early for rate limit!")
+                sys.exit(-1)
+
             Output.warn(
                 f"Sleeping for {Output.bright( sleep_time_mins + ' minutes')} "
                 "to prevent rate limit exhaustion!")
 
-            if self.no_sleep:
-                Output.warn("Skipping sleep due to no_sleep being set to True (exiting instead)")
-                sys.exit(-1)
-
             time.sleep(sleep_time + 1)
 
     def __process_run_log(self, log_content: bytes, run_info: dict):