Add Optional Argument to Disable Sleeping

Author: UNC1739

gato/attack/attack.py

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

gato/cli/cli.py

@@ -49,6 +49,12 @@ 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

gato/enumerate/enumerate.py

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

gato/github/api.py

@@ -28,7 +28,7 @@ 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"):
+                 github_url: str = "https://api.github.com", no_sleep: bool = False):
         """Initialize the API abstraction layer to interact with the GitHub
         REST API.
 
@@ -50,6 +50,8 @@ def __init__(self, pat: str, version: str = "2022-11-28",
             'Authorization': f'Bearer {pat}',
             'X-GitHub-Api-Version': version
         }
+        self.no_sleep = no_sleep
+
         if not github_url:
             self.github_url = "https://api.github.com"
         else:
@@ -103,6 +105,10 @@ def __check_rate_limit(self, headers):
                 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):