feat: make GitHub searches resilient

Author: techgaun

README.md

@@ -60,11 +60,18 @@ GH_USER=techgaun GH_PWD=<mypass> github-dorks -u dev-nepal     # search as authe
 GH_TOKEN=<github_token> github-dorks -u dev-nepal              # search using auth token
 
 GH_URL=https://github.example.com github-dorks -u dev-nepal    # search a GitHub Enterprise instance
+
+github-dorks -r techgaun/github-dorks --max-retries 5         # retry recoverable failures up to five times
 ```
 
 The legacy `python github-dork.py ...` invocation remains available for
 compatibility. The package can also run as `python -m github_dorks ...`.
 
+Each scan ends with a summary of queries, matches, failures, retries, and
+elapsed time. The command exits with status `0` after a complete scan, `2` when
+one or more queries failed, and `1` for fatal configuration, file, or
+authentication errors.
+
 ### Development
 
 Run the dependency-free unit test suite with:
@@ -77,9 +84,9 @@ The CI test matrix covers Python 3.10 through 3.13.
 
 ### Limitations
 
-- Authenticated requests get a higher rate limit. But, since this tool waits for the api rate limit to be reset (which is usually less than a minute), it can be slightly slow.
+- Authenticated requests receive higher rate limits. Searches may pause until
+  GitHub resets the search limit.
 - Search results can be printed to the terminal or written as CSV.
-- ~~Handle rate limit and retry. PR welcome~~
 
 ### Contribution
 

github-dork.py

@@ -5,4 +5,4 @@
 
 
 if __name__ == '__main__':
-    main()
+    raise SystemExit(main())

github_dorks/__main__.py

@@ -2,4 +2,4 @@
 
 
 if __name__ == '__main__':
-    main()
+    raise SystemExit(main())

github_dorks/cli.py

@@ -1,225 +1,101 @@
-import github3 as github
-import os
+"""Command-line interface for github-dorks."""
+
 import argparse
-import csv
+import os
+import sys
 import time
+
 import feedparser
-from copy import copy
-from contextlib import nullcontext
-from sys import stderr, prefix
+import github3 as github
 
 from github_dorks import __version__
+from github_dorks.search import create_client, search
 
-gh_user = os.getenv('GH_USER', None)
-gh_pass = os.getenv('GH_PWD', None)
-gh_token = os.getenv('GH_TOKEN', None)
-gh_url = os.getenv('GH_URL', None)
 
-if gh_url is None:
-    gh = github.GitHub(username=gh_user, password=gh_pass, token=gh_token)
-else:
-    gh = github.GitHubEnterprise(
-        url=gh_url, username=gh_user, password=gh_pass, token=gh_token)
+def monitor(gh_dorks_file=None, feed_token=None, refresh_time=60):
+    github_user = os.getenv('GH_USER')
+    if github_user is None:
+        raise ValueError('GH_USER is required for monitoring')
 
-
-def search_wrapper(gen):
+    print('Monitoring merged pull requests for new scans.')
+    seen_items = set()
+    private_feed = f'https://github.com/{github_user}.private.atom?token={feed_token}'
     while True:
-        gen_back = copy(gen)
-        try:
-            yield next(gen)
-        except StopIteration:
-            return
-        except github.exceptions.ForbiddenError:
-            search_rate_limit = gh.rate_limit()['resources']['search']
-            # limit_remaining = search_rate_limit['remaining']
-            reset_time = search_rate_limit['reset']
-            current_time = int(time.time())
-            sleep_time = reset_time - current_time + 1
-            stderr.write(
-                'GitHub Search API rate limit reached. Sleeping for %d seconds.\n\n'
-                % (sleep_time))
-            time.sleep(sleep_time)
-            yield next(gen_back)
-        except Exception as e:
-            raise e
-
-
-def metasearch(repo_to_search=None,
-               user_to_search=None,
-               gh_dorks_file=None,
-               active_monit=None,
-               output_filename=None,
-               refresh_time=60):
-    if active_monit is None:
-        search(repo_to_search, user_to_search, gh_dorks_file, active_monit, output_filename)
-    else:
-        monit(gh_dorks_file, active_monit, refresh_time)
-
-
-def monit(gh_dorks_file=None, active_monit=None, refresh_time=60):
-    if gh_user is None:
-        raise Exception('Error, env Github user variable needed')
-    else:
-        print(
-            'Monitoring user private feed searching new code to be dorked.' +
-            'Every new merged pull request trigger user scan.'
-        )
-        print('-----')
-        items_history = list()
-        gh_private_feed = "https://github.com/{}.private.atom?token={}".format(
-            gh_user, active_monit)
-        while True:
-            feed = feedparser.parse(gh_private_feed)
-            for i in feed['items']:
-                if 'merged pull' in i['title']:
-                    if i['title'] not in items_history:
-                        search(
-                            user_to_search=i['author_detail']['name'],
-                            gh_dorks_file=gh_dorks_file)
-                        items_history.append(i['title'])
-            print('Waiting for new items...')
-            time.sleep(refresh_time)
-
-
-def search(repo_to_search=None,
-           user_to_search=None,
-           gh_dorks_file=None,
-           active_monit=None,
-           output_filename=None):
-
-    if gh_dorks_file is None:
-        for path_prefix in ['.', os.path.join(prefix, 'github-dorks/')]:
-            filename = os.path.join(path_prefix, 'github-dorks.txt')
-            if os.path.isfile(filename):
-                gh_dorks_file = filename
-                break
-
-    if gh_dorks_file is None or not os.path.isfile(gh_dorks_file):
-        raise Exception('Error, the dorks file path is not valid')
-    if user_to_search:
-        print("Scanning User: ", user_to_search)
-    if repo_to_search:
-        print("Scanning Repo: ", repo_to_search)
-    found = False
-
-    output_context = (
-        open(output_filename, 'w', newline='', encoding='utf-8')
-        if output_filename else nullcontext(None)
-    )
-
-    with open(gh_dorks_file, 'r', encoding='utf-8') as dork_file, output_context as output_file:
-        # Write CSV Header
-        csv_writer = None
-        if output_file:
-            csv_writer = csv.writer(output_file)
-            csv_writer.writerow([
-                'Issue Type (Dork)', 'Text Matches', 'File Path',
-                'Score/Relevance', 'URL of File'
-            ])
-        for dork in dork_file:
-            dork = dork.strip()
-            if not dork or dork[0] in '#;':
-                continue
-            addendum = ''
-            if repo_to_search:
-                addendum = ' repo:' + repo_to_search
-            elif user_to_search:
-                addendum = ' user:' + user_to_search
-
-            dork = dork + addendum
-            search_results = search_wrapper(gh.search_code(dork))
-            try:
-                for search_result in search_results:
-                    found = True
-                    fmt_args = {
-                        'dork': dork,
-                        'text_matches': search_result.text_matches,
-                        'path': search_result.path,
-                        'score': search_result.score,
-                        'url': search_result.html_url
-                    }
-
-                    # Either write to file or print output
-                    if csv_writer:
-                        csv_writer.writerow([
-                            fmt_args['dork'], fmt_args['text_matches'],
-                            fmt_args['path'], fmt_args['score'], fmt_args['url']
-                        ])
-                    else:
-                        result = '\n'.join([
-                            'Found result for {dork}',
-                            'Text matches: {text_matches}', 'File path: {path}',
-                            'Score/Relevance: {score}', 'URL of File: {url}', ''
-                        ]).format(**fmt_args)
-                        print(result)
-
-            except github.exceptions.GitHubError as e:
-                print('GitHubError encountered on search of dork: ' + dork)
-                print(e)
-                return
-            except Exception as e:
-                print(e)
-                print('Error encountered on search of dork: ' + dork)
-
-    if not found:
-        print('No results for your dork search' + addendum + '. Hurray!')
-
-
-def main():
+        feed = feedparser.parse(private_feed)
+        for item in feed['items']:
+            if 'merged pull' in item['title'] and item['title'] not in seen_items:
+                search(
+                    user_to_search=item['author_detail']['name'],
+                    gh_dorks_file=gh_dorks_file,
+                )
+                seen_items.add(item['title'])
+        print('Waiting for new items...')
+        time.sleep(refresh_time)
+
+
+def build_parser():
     parser = argparse.ArgumentParser(
         prog='github-dorks',
         description='Search GitHub for sensitive data patterns',
-        epilog='Use responsibly. Only scan repositories you are authorized to assess.')
-
+        epilog='Use responsibly. Only scan repositories you are authorized to assess.',
+    )
     parser.add_argument(
-        '-v', '--version', action='version', version='%(prog)s ' + __version__)
-
+        '-v', '--version', action='version', version='%(prog)s ' + __version__
+    )
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument(
-        '-u',
-        '--user',
-        dest='user_to_search',
-        action='store',
-        help='GitHub user/org to search within. Eg: techgaun')
-
+        '-u', '--user', dest='user_to_search',
+        help='GitHub user/org to search within. Eg: techgaun',
+    )
     group.add_argument(
-        '-r',
-        '--repo',
-        dest='repo_to_search',
-        action='store',
-        help='GitHub repo to search within. Eg: techgaun/github-dorks')
-
-    parser.add_argument(
-        '-d',
-        '--dork',
-        dest='gh_dorks_file',
-        action='store',
-        help='GitHub dorks file. Eg: github-dorks.txt')
-
+        '-r', '--repo', dest='repo_to_search',
+        help='GitHub repo to search within. Eg: techgaun/github-dorks',
+    )
     group.add_argument(
-        '-m',
-        '--monit',
-        dest='active_monit',
-        action='store',
-        help='Monitors GitHub user private feed with feed token'
+        '-m', '--monit', dest='active_monit',
+        help='Monitor the GitHub user private feed with this feed token',
     )
-
     parser.add_argument(
-        '-o',
-        '--outputFile',
-        dest='output_filename',
-        action='store',
-        help='CSV File to write results to. This overwrites the file provided! Eg: out.csv'
+        '-d', '--dork', dest='gh_dorks_file',
+        help='GitHub dorks file. Eg: github-dorks.txt',
     )
-
-    args = parser.parse_args()
-    metasearch(
-        repo_to_search=args.repo_to_search,
-        user_to_search=args.user_to_search,
-        gh_dorks_file=args.gh_dorks_file,
-        active_monit=args.active_monit,
-        output_filename=args.output_filename)
+    parser.add_argument(
+        '-o', '--outputFile', dest='output_filename',
+        help='CSV file to write results to. This overwrites the provided file.',
+    )
+    parser.add_argument(
+        '--max-retries', type=int, default=3,
+        help='Maximum retries per query for recoverable API failures (default: 3)',
+    )
+    return parser
 
 
-if __name__ == '__main__':
-    main()
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+    if args.max_retries < 0:
+        parser.error('--max-retries must be zero or greater')
+    try:
+        if args.active_monit:
+            monitor(args.gh_dorks_file, args.active_monit)
+            return 0
+        stats = search(
+            repo_to_search=args.repo_to_search,
+            user_to_search=args.user_to_search,
+            gh_dorks_file=args.gh_dorks_file,
+            output_filename=args.output_filename,
+            client=create_client(),
+            max_retries=args.max_retries,
+        )
+        return stats.exit_code
+    except (OSError, ValueError) as error:
+        print(f'Error: {error}', file=sys.stderr)
+        return 1
+    except Exception as error:
+        authentication_error = getattr(
+            github.exceptions, 'AuthenticationFailed', ()
+        )
+        if authentication_error and isinstance(error, authentication_error):
+            print(f'Error: {error}', file=sys.stderr)
+            return 1
+        raise