From 5eb1f83f0c02a47f1b61c0cded8bec81bbf4608a Mon Sep 17 00:00:00 2001
From: Denis Loginov <dinvlad@gmail.com>
Date: Tue, 13 Oct 2020 20:29:07 -0400
Subject: [PATCH] Add --file-pattern argument

---
 README.md                 | 13 ++++++++++---
 find_gcp_keys/__main__.py | 23 ++++++++++++++---------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index a74347a..d95c8a2 100644
--- a/README.md
+++ b/README.md
@@ -18,13 +18,20 @@ pip3 install find-gcp-keys
 
 ## Usage
 
-As a command-line utility:
+### Command line
 
 ```
-find-gcp-keys <dir_path> [--no-validate/-n]
+find-gcp-keys <dir_path> [--no-validate/-n] [--file-pattern <regex>]
 ```
 
-As a library:
+Note that by default, the CLI only searches for the JSON key files
+matching a particular pattern (`<project-id>-<key-id>.json`). You can
+override this behavior, e.g. to search for _all_ JSON files:
+```
+find-gcp-keys <dir_path> -p '.*\.json'
+```
+
+### Library:
 
 ```py
 from find_gcp_keys import find_key_paths, find_valid_keys, is_valid_key
diff --git a/find_gcp_keys/__main__.py b/find_gcp_keys/__main__.py
index 3dbb0af..3bfc6b6 100755
--- a/find_gcp_keys/__main__.py
+++ b/find_gcp_keys/__main__.py
@@ -27,20 +27,25 @@ def parse_args():
         '--no-validate', '-n', action='store_true',
         help='Directory path to search recursively',
     )
+    parser.add_argument(
+        '--file-pattern', '-p', default=PROJECT_JSON_PATTERN,
+        help='Pattern to match JSON file names against',
+    )
     return parser.parse_args()
 
 
 # For requirements on GCP project IDs, see
 # https://cloud.google.com/resource-manager/docs/creating-managing-projects
 PROJECT_PATTERN = r"[a-z][a-z0-9\-]{4,28}[a-z0-9]"
-FILE_PATTERN = re.compile(PROJECT_PATTERN + r"-[0-9a-f]{12}\.json")
+PROJECT_JSON_PATTERN = PROJECT_PATTERN + r"-[0-9a-f]{12}\.json"
 
 
-def find_key_paths(dir_path: str):
-    """ Finds files whose name matches the JSON SA key pattern """
+def find_key_paths(dir_path: str, file_pattern: str):
+    """ Finds files whose name matches `file_pattern` """
+    file_re = re.compile(file_pattern)
     for dirpath, _, files in os.walk(dir_path):
         for file in files:
-            if FILE_PATTERN.match(file):
+            if file_re.match(file):
                 yield os.path.join(dirpath, file)
 
 
@@ -52,13 +57,13 @@ def is_valid_key(file_path: str):
         )
         credentials.refresh(google.auth.transport.requests.Request())
         return True
-    except (ValueError, google.auth.exceptions.RefreshError):
+    except (AttributeError, ValueError, google.auth.exceptions.RefreshError):
         return False
 
 
-def find_valid_keys(dir_path: str):
+def find_valid_keys(dir_path: str, file_pattern: str):
     """ Recursively walks `dir_path` and finds valid GCP SA keys """
-    for path in find_key_paths(dir_path):
+    for path in find_key_paths(dir_path, file_pattern):
         if is_valid_key(path):
             yield path
 
@@ -68,12 +73,12 @@ def main():
     args = parse_args()
 
     if args.no_validate:
-        for path in find_key_paths(args.dir_path):
+        for path in find_key_paths(args.dir_path, args.file_pattern):
             print(path)
         sys.exit(0)
 
     found = False
-    for path in find_valid_keys(args.dir_path):
+    for path in find_valid_keys(args.dir_path, args.file_pattern):
         print(path, file=sys.stderr)
         found = True