diff --git a/gato/attack/attack.py b/gato/attack/attack.py index 4f82357..4707519 100644 --- a/gato/attack/attack.py +++ b/gato/attack/attack.py @@ -38,7 +38,7 @@ def __init__( author_name: str = None, timeout: int = 30, github_url: str = None, - no_sleep: bool = False + no_sleep: bool = False ): self.api = Api( @@ -46,7 +46,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 diff --git a/gato/cli/cli.py b/gato/cli/cli.py index df7c3b3..ef31f60 100644 --- a/gato/cli/cli.py +++ b/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. diff --git a/gato/enumerate/enumerate.py b/gato/enumerate/enumerate.py index 12ebf2d..a33139f 100644 --- a/gato/enumerate/enumerate.py +++ b/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 diff --git a/gato/github/api.py b/gato/github/api.py index 5a81ab1..f5232eb 100644 --- a/gato/github/api.py +++ b/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):