From 6e54e55c317c34d8671f0ca846e3e3b23f9032ee Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 4 Oct 2020 19:42:44 +0300
Subject: [PATCH] Add agent.force_git_ssh_port to control https to ssh link
 conversion for non standard ssh port

---
 docs/trains.conf                                   |  5 +++++
 trains_agent/backend_api/config/default/agent.conf |  2 ++
 trains_agent/helper/repo.py                        | 11 ++++++-----
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/docs/trains.conf b/docs/trains.conf
index 5032e0f..42e3c73 100644
--- a/docs/trains.conf
+++ b/docs/trains.conf
@@ -17,9 +17,14 @@ agent {
     # leave blank for GIT SSH credentials (set force_git_ssh_protocol=true to force SSH protocol)
     git_user=""
     git_pass=""
+    # Limit credentials to a single domain, for example: github.com,
+    # all other domains will use public access (no user/pass). Default: always send user/pass for any VCS domain
+    git_domain=""
 
     # Force GIT protocol to use SSH regardless of the git url (Assumes GIT user/pass are blank)
     force_git_ssh_protocol: false
+    # Force a specific SSH port when converting http to ssh links (the domain is kept the same)
+    # force_git_ssh_port: ""
 
     # unique name of this worker, if None, created based on hostname:process_id
     # Overridden with os environment: TRAINS_WORKER_NAME
diff --git a/trains_agent/backend_api/config/default/agent.conf b/trains_agent/backend_api/config/default/agent.conf
index 2888844..e82a40b 100644
--- a/trains_agent/backend_api/config/default/agent.conf
+++ b/trains_agent/backend_api/config/default/agent.conf
@@ -16,6 +16,8 @@
 
     # Force GIT protocol to use SSH regardless of the git url (Assumes GIT user/pass are blank)
     force_git_ssh_protocol: false
+    # Force a specific SSH port when converting http to ssh links (the domain is kept the same)
+    # force_git_ssh_port: 0
 
     # Set the python version to use when creating the virtual environment and launching the experiment
     # Example values: "/usr/bin/python3" or "/usr/local/bin/python3.6"
diff --git a/trains_agent/helper/repo.py b/trains_agent/helper/repo.py
index 4dfea20..00fd63f 100644
--- a/trains_agent/helper/repo.py
+++ b/trains_agent/helper/repo.py
@@ -5,7 +5,7 @@ import subprocess
 from distutils.spawn import find_executable
 from hashlib import md5
 from os import environ, getenv
-from typing import Text, Sequence, Mapping, Iterable, TypeVar, Callable, Tuple
+from typing import Text, Sequence, Mapping, Iterable, TypeVar, Callable, Tuple, Optional
 
 import attr
 from furl import furl
@@ -254,8 +254,8 @@ class VCS(object):
         return url
 
     @classmethod
-    def replace_http_url(cls, url):
-        # type: (Text) -> Text
+    def replace_http_url(cls, url, port=None):
+        # type: (Text, Optional[int]) -> Text
         """
         Replace HTTPS URL with SSH URL when applicable
         """
@@ -267,7 +267,7 @@ class VCS(object):
             # make sure there is no port in the final url (safe_furl support)
             # the original port was an https port, and we do not know if there is a different ssh port,
             # so we have to clear the original port specified (https) and use the default ssh schema port.
-            parsed_url.port = None
+            parsed_url.port = port or None
             url = parsed_url.url
         return url
 
@@ -279,7 +279,8 @@ class VCS(object):
         if self.session.config.get('agent.force_git_ssh_protocol', None) and self.url:
             parsed_url = furl(self.url)
             if parsed_url.scheme == "https":
-                new_url = self.replace_http_url(self.url)
+                new_url = self.replace_http_url(
+                    self.url, port=self.session.config.get('agent.force_git_ssh_port', None))
                 if new_url != self.url:
                     print("Using SSH credentials - replacing https url '{}' with ssh url '{}'".format(
                         self.url, new_url))