From 9456e493ac6d6495310ee084db906f9cdca8218c Mon Sep 17 00:00:00 2001
From: Simon Gasse <sgasse@users.noreply.github.com>
Date: Tue, 27 Jul 2021 14:00:33 +0200
Subject: [PATCH] Enable rewriting SSH URLs

ClearML Agent allows to force git cloning via SSH and also has a
setting to force a username. The relevant settings are:
agent.force_git_ssh_protocol: true
agent.force_git_ssh_user: "git"

However, forcing a specific username or port only worked so far if the
agent translated either from https->ssh or from ssh->https. A given
ssh URL was not rewritten.

This commit adds a helper function and includes it in `_set_ssh_url`
to allow rewriting ssh URLs with the username and/or port given in the
config `agent.force_git_ssh_user`.
If neither username nor port are forced in the config, the URL is not
touched.

This is somewhat related to issue #42.
Note that rewriting https->https is not covered in this commit.
---
 clearml_agent/helper/repo.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/clearml_agent/helper/repo.py b/clearml_agent/helper/repo.py
index 551c9f3..0bab3f2 100644
--- a/clearml_agent/helper/repo.py
+++ b/clearml_agent/helper/repo.py
@@ -274,6 +274,18 @@ class VCS(object):
             url = parsed_url.url
         return url
 
+    @classmethod
+    def rewrite_ssh_url(cls, url, port=None, username=None):
+        # type: (Text, Optional[int], Optional[str]) -> Text
+        """
+        Rewrite SSH URL with custom port and username
+        """
+        parsed_url = furl(url)
+        if parsed_url.scheme == "ssh":
+            parsed_url.username = username or "git"
+            parsed_url.port = port or None
+            return parsed_url.url
+
     def _set_ssh_url(self):
         """
         Replace instance URL with SSH substitution result and report to log.
@@ -297,6 +309,20 @@ class VCS(object):
                         self.url, new_url))
                     self.url = new_url
                 return
+            # rewrite ssh URLs only if either ssh port or ssh user are forced in config
+            if parsed_url.scheme == "ssh" and (
+                self.session.config.get('agent.force_git_ssh_port', None) or
+                self.session.config.get('agent.force_git_ssh_user', None)
+            ):
+                new_url = self.rewrite_ssh_url(
+                    self.url,
+                    port=self.session.config.get('agent.force_git_ssh_port', None),
+                    username=self.session.config.get('agent.force_git_ssh_user', None)
+                )
+                if new_url != self.url:
+                    print("Using SSH credentials - ssh url '{}' with ssh url '{}'".format(
+                        self.url, new_url))
+                    self.url = new_url
 
         if not self.session.config.agent.translate_ssh:
             return