Better

docker/docs/builder/Dockerfile

@@ -30,6 +30,7 @@ RUN apt-get update \
             nodejs \
             npm \
             openjdk-11-jdk \
+            ssh-client \
     && pip --no-cache-dir install scipy \
     && apt-get autoremove --yes \
     && apt-get clean \

tests/ci/docs_release.py

@@ -12,6 +12,7 @@ from report import create_test_html_report
 from s3_helper import S3Helper
 from pr_info import PRInfo
 from get_robot_token import get_best_robot_token, get_parameter_from_ssm
+from ssh import SSHKey
 
 NAME = "Docs Release (actions)"
 
@@ -109,7 +110,7 @@ if __name__ == "__main__":
 
     run_log_path = os.path.join(test_output, 'runlog.log')
 
-    with open(run_log_path, 'w', encoding='utf-8') as log:
+    with open(run_log_path, 'w', encoding='utf-8') as log, SSHKey("robot-clickhouse-ssh"):
         with subprocess.Popen(cmd, shell=True, stderr=log, stdout=log) as process:
             retcode = process.wait()
             if retcode == 0:
@@ -149,5 +150,4 @@ if __name__ == "__main__":
 
     report_url = upload_results(s3_helper, pr_info.number, pr_info.sha, lines, additional_files)
     print("::notice ::Report url: {report_url}")
-    commit = get_commit(gh, pr_info.sha)
     commit.create_status(context=NAME, description=description, state=status, target_url=report_url)

tests/ci/ssh.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+
+import shutil
+import os
+import subprocess
+import tempfile
+import logging
+import signal
+from get_robot_token import get_parameter_from_ssm
+
+
+class SSHAgent:
+    def __init__(self):
+        self._env = {}
+        self._env_backup = {}
+        self._keys = {}
+        self.start()
+
+    @property
+    def pid(self):
+        return int(self._env["SSH_AGENT_PID"])
+
+    def start(self):
+        if shutil.which("ssh-agent") is None:
+            raise Exception("ssh-agent binary is not available")
+
+        self._env_backup["SSH_AUTH_SOCK"] = os.environ.get("SSH_AUTH_SOCK")
+        self._env_backup["SSH_OPTIONS"] = os.environ.get("SSH_OPTIONS")
+
+        # set ENV from stdout of ssh-agent
+        for line in self._run(['ssh-agent']).splitlines():
+            name, _, value = line.partition(b"=")
+            if _ == b"=":
+                value = value.split(b";", 1)[0]
+                self._env[name.decode()] = value.decode()
+                os.environ[name.decode()] = value.decode()
+
+        ssh_options = "," + os.environ["SSH_OPTIONS"] if os.environ.get("SSH_OPTIONS") else ""
+        os.environ["SSH_OPTIONS"] = f"{ssh_options}UserKnownHostsFile=/dev/null,StrictHostKeyChecking=no"
+
+    def add(self, key):
+        key_pub = self._key_pub(key)
+
+        if key_pub in self._keys:
+            self._keys[key_pub] += 1
+        else:
+            self._run(["ssh-add", "-"], stdin=key.encode())
+            self._keys[key_pub] = 1
+
+        return key_pub
+
+    def remove(self, key_pub):
+        if key_pub not in self._keys:
+            raise Exception(f"Private key not found, public part: {key_pub}")
+
+        if self._keys[key_pub] > 1:
+            self._keys[key_pub] -= 1
+        else:
+            with tempfile.NamedTemporaryFile() as f:
+                f.write(key_pub)
+                f.flush()
+                self._run(["ssh-add", "-d", f.name])
+            self._keys.pop(key_pub)
+
+    def print_keys(self):
+        keys = self._run(["ssh-add", "-l"]).splitlines()
+        if keys:
+            logging.info("ssh-agent keys:")
+            for key in keys:
+                logging.info("%s", key)
+        else:
+            logging.info("ssh-agent (pid %d) is empty", self.pid)
+
+    def kill(self):
+        for k, v in self._env.items():
+            os.environ.pop(k, None)
+
+        for k, v in self._env_backup.items():
+            if v is not None:
+                os.environ[k] = v
+
+        os.kill(self.pid, signal.SIGTERM)
+
+    def _key_pub(self, key):
+        with tempfile.NamedTemporaryFile() as f:
+            f.write(key.encode())
+            f.flush()
+            return self._run(["ssh-keygen", "-y", "-f", f.name])
+
+    @staticmethod
+    def _run(cmd, stdin=None):
+        shell = isinstance(cmd, str)
+        with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE if stdin else None, shell=shell) as p:
+            stdout, stderr = p.communicate(stdin)
+
+            if stdout.strip().decode() == "The agent has no identities.":
+                return ""
+
+            if p.returncode:
+                message = stderr.strip() + b"\n" + stdout.strip()
+                raise Exception(message.strip().decode())
+
+            return stdout
+
+class SSHKey:
+    def __init__(self, key_name):
+        self.key = get_parameter_from_ssm(key_name)
+        self._key_pub = None
+        self._ssh_agent = SSHAgent()
+
+    def __enter__(self):
+        self._key_pub = self._ssh_agent.add(self.key)
+        self._ssh_agent.print_keys()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._ssh_agent.remove(self._key_pub)
+        self._ssh_agent.print_keys()