added support for auth_method in push() and authenticate()

Author: ivanharvard

lib50/_api.py

@@ -39,7 +39,7 @@
 DEFAULT_FILE_LIMIT = 10000
 
 
-def push(tool, slug, config_loader, repo=None, data=None, prompt=lambda question, included, excluded: True, file_limit=DEFAULT_FILE_LIMIT):
+def push(tool, slug, config_loader, repo=None, data=None, prompt=lambda question, included, excluded: True, file_limit=DEFAULT_FILE_LIMIT, auth_method=None):
     """
     Pushes to Github in name of a tool.
     What should be pushed is configured by the tool and its configuration in the .cs50.yml file identified by the slug.
@@ -61,6 +61,10 @@ def push(tool, slug, config_loader, repo=None, data=None, prompt=lambda question
     :type prompt: lambda str, list, list => bool, optional
     :param file_limit: maximum number of files to be matched by any globbing pattern.
     :type file_limit: int
+    :param auth_method: The authentication method to use.  Accepts `"https"` or `"ssh"`. \
+                        If any other value is provided, attempts SSH \
+                        authentication first and fall back to HTTPS if SSH fails.
+    :type auth_method: str
     :return: GitHub username and the commit hash
     :type: tuple(str, str)
 
@@ -89,7 +93,7 @@ def push(tool, slug, config_loader, repo=None, data=None, prompt=lambda question
     remote, (honesty, included, excluded) = connect(slug, config_loader, file_limit=DEFAULT_FILE_LIMIT)
 
     # Authenticate the user with GitHub, and prepare the submission
-    with authenticate(remote["org"], repo=repo) as user, prepare(tool, slug, user, included):
+    with authenticate(remote["org"], repo=repo, auth_method=auth_method) as user, prepare(tool, slug, user, included):
 
         # Show any prompt if specified
         if prompt(honesty, included, excluded):
@@ -465,7 +469,7 @@ def batch_files(files, size=100):
                 files_list = list(files)
                 for i in range(0, len(files_list), size):
                     yield files_list[i:i + size]
-            
+
             for batch in batch_files(included):
                 quoted_files = ' '.join(shlex.quote(f) for f in batch)
                 run(git(f"add -f -- {quoted_files}"))
@@ -962,7 +966,7 @@ def run(command, quiet=False, timeout=None):
                 ProgressBar.stop_all()
                 passphrase = _prompt_password("Enter passphrase for SSH key: ")
                 child.sendline(passphrase)
-        
+
                 # Get the full output by reading until EOF
                 full_output = child.before + child.after + child.read()
                 command_output = full_output.strip().replace("\r\n", "\n")

lib50/authentication.py

@@ -32,14 +32,18 @@ class User:
                     init=False)
 
 @contextlib.contextmanager
-def authenticate(org, repo=None):
+def authenticate(org, repo=None, auth_method=None):
     """
-    A contextmanager that authenticates a user with GitHub via SSH if possible, otherwise via HTTPS.
+    A contextmanager that authenticates a user with GitHub.
 
     :param org: GitHub organisation to authenticate with
     :type org: str
     :param repo: GitHub repo (part of the org) to authenticate with. Default is the user's GitHub login.
     :type repo: str, optional
+    :param auth_method: The authentication method to use.  Accepts `"https"` or `"ssh"`. \
+                        If any other value is provided, attempts SSH \
+                        authentication first and fall back to HTTPS if SSH fails.
+    :type auth_method: str, optional
     :return: an authenticated user
     :type: lib50.User
 
@@ -51,21 +55,34 @@ def authenticate(org, repo=None):
             print(user.name)
 
     """
-    with api.ProgressBar(_("Authenticating")) as progress_bar:
-        # Both authentication methods can require user input, best stop the bar
-        progress_bar.stop()
+    def try_https(org, repo):
+        with _authenticate_https(org, repo=repo) as user:
+            return user
 
-        # Try auth through SSH
+    def try_ssh(org, repo):
         user = _authenticate_ssh(org, repo=repo)
-
-        # SSH auth failed, fallback to HTTPS
         if user is None:
-            with _authenticate_https(org, repo=repo) as user:
-                yield user
-        # yield SSH user
-        else:
-            yield user
+            raise ConnectionError
+        return user
+
+    # Showcase the type of authentication based on input
+    method_label = f" ({auth_method.upper()})" if auth_method in ("https", "ssh") else ""
+    with api.ProgressBar(_("Authenticating{}").format(method_label)) as progress_bar:
+        # Both authentication methods can require user input, best stop the bar
+        progress_bar.stop()
 
+        match auth_method:
+            case "https":
+                yield try_https(org, repo)
+            case "ssh":
+                yield try_ssh(org, repo)
+            case _:
+                # Try auth through SSH
+                try:
+                    yield try_ssh(org, repo)
+                except ConnectionError:
+                    # SSH auth failed, fallback to HTTPS
+                    yield try_https(org, repo)
 
 def logout():
     """
@@ -104,7 +121,7 @@ def run_authenticated(user, command, quiet=False, timeout=None):
                 # Try to extract the conflicting branch prefix from the error message
                 # Pattern: 'refs/heads/cs50/problems/2025/x' exists
                 branch_prefix_match = re.search(r"'refs/heads/([^']+)' exists", command_output)
-                
+
                 if branch_prefix_match:
                     conflicting_prefix = branch_prefix_match.group(1)
                     error_msg = _("Looks like you're trying to push to a branch that conflicts with an existing one in the repository.\n"
@@ -195,7 +212,7 @@ class State(enum.Enum):
         else:
             if not os.environ.get("CODESPACES"):
                 _show_gh_changes_warning()
-                
+
             return None
     finally:
         child.close()