From f07fcbece1eaf77931c0419991e138fa4e89c779 Mon Sep 17 00:00:00 2001
From: "William D. Jones" <thor0505@comcast.net>
Date: Wed, 26 Aug 2020 18:49:49 -0400
Subject: [PATCH] build.run: implement SSH remote builds using Paramiko.

---
 nmigen/build/run.py | 112 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 108 insertions(+), 4 deletions(-)

diff --git a/nmigen/build/run.py b/nmigen/build/run.py
index 82c27ad..507237a 100644
--- a/nmigen/build/run.py
+++ b/nmigen/build/run.py
@@ -7,9 +7,11 @@ import subprocess
 import tempfile
 import zipfile
 import hashlib
+import pathlib
 
 
-__all__ = ["BuildPlan", "BuildProducts", "LocalBuildProducts"]
+__all__ = ["BuildPlan", "BuildProducts", "LocalBuildProducts", "RemoteSSHBuildProducts"]
+
 
 
 class BuildPlan:
@@ -74,9 +76,10 @@ class BuildPlan:
             os.chdir(root)
 
             for filename, content in self.files.items():
-                filename = os.path.normpath(filename)
-                # Just to make sure we don't accidentally overwrite anything outside of build root.
-                assert not filename.startswith("..")
+                filename = pathlib.Path(filename)
+                # Forbid parent directory components completely to avoid the possibility
+                # of writing outside the build root.
+                assert ".." not in filename.parts
                 dirname = os.path.dirname(filename)
                 if dirname:
                     os.makedirs(dirname, exist_ok=True)
@@ -99,6 +102,82 @@ class BuildPlan:
         finally:
             os.chdir(cwd)
 
+    def execute_remote_ssh(self, *, connect_to = {}, root, run_script=True):
+        """
+        Execute build plan using the remote SSH strategy. Files from the build
+        plan are transferred via SFTP to the directory ``root`` on a  remote
+        server. If ``run_script`` is ``True``, the ``paramiko`` SSH client will
+        then run ``{script}.sh``. ``root`` can either be an absolute or
+        relative (to the login directory) path.
+
+        ``connect_to`` is a dictionary that holds all input arguments to
+        ``paramiko``'s ``SSHClient.connect``
+        (`documentation <http://docs.paramiko.org/en/stable/api/client.html#paramiko.client.SSHClient.connect>`_).
+        At a minimum, the ``hostname`` input argument must be supplied in this
+        dictionary as the remote server.
+
+        Returns :class:`RemoteSSHBuildProducts`.
+        """
+        from paramiko import SSHClient
+
+        with SSHClient() as client:
+            client.load_system_host_keys()
+            client.connect(**connect_to)
+
+            with client.open_sftp() as sftp:
+                def mkdir_exist_ok(path):
+                    try:
+                        sftp.mkdir(str(path))
+                    except IOError as e:
+                        # mkdir fails if directory exists. This is fine in nmigen.build.
+                        # Reraise errors containing e.errno info.
+                        if e.errno:
+                            raise e
+
+                def mkdirs(path):
+                    # Iteratively create parent directories of a file by iterating over all
+                    # parents except for the root ("."). Slicing the parents results in
+                    # TypeError, so skip over the root ("."); this also handles files
+                    # already in the root directory.
+                    for parent in reversed(path.parents):
+                        if parent == pathlib.PurePosixPath("."):
+                            continue
+                        else:
+                            mkdir_exist_ok(parent)
+
+                mkdir_exist_ok(root)
+
+                sftp.chdir(root)
+                for filename, content in self.files.items():
+                    filename = pathlib.PurePosixPath(filename)
+                    assert ".." not in filename.parts
+
+                    mkdirs(filename)
+
+                    mode = "wt" if isinstance(content, str) else "wb"
+                    with sftp.file(str(filename), mode) as f:
+                        # "b/t" modifier ignored in SFTP.
+                        if mode == "wt":
+                            f.write(content.encode("utf-8"))
+                        else:
+                            f.write(content)
+
+            if run_script:
+                transport = client.get_transport()
+                channel = transport.open_session()
+                channel.set_combine_stderr(True)
+
+                cmd = "if [ -f ~/.profile ]; then . ~/.profile; fi && cd {} && sh {}.sh".format(root, self.script)
+                channel.exec_command(cmd)
+
+                # Show the output from the server while products are built.
+                buf = channel.recv(1024)
+                while buf:
+                    print(buf.decode("utf-8"), end="")
+                    buf = channel.recv(1024)
+
+        return RemoteSSHBuildProducts(connect_to, root)
+
     def execute(self):
         """
         Execute build plan using the default strategy. Use one of the ``execute_*`` methods
@@ -162,3 +241,28 @@ class LocalBuildProducts(BuildProducts):
         super().get(filename, mode)
         with open(os.path.join(self.__root, filename), "r" + mode) as f:
             return f.read()
+
+
+class RemoteSSHBuildProducts(BuildProducts):
+    def __init__(self, connect_to, root):
+        self.__connect_to = connect_to
+        self.__root = root
+
+    def get(self, filename, mode="b"):
+        super().get(filename, mode)
+
+        from paramiko import SSHClient
+
+        with SSHClient() as client:
+            client.load_system_host_keys()
+            client.connect(**self.__connect_to)
+
+            with client.open_sftp() as sftp:
+                sftp.chdir(self.__root)
+
+                with sftp.file(filename, "r" + mode) as f:
+                    # "b/t" modifier ignored in SFTP.
+                    if mode == "t":
+                        return f.read().decode("utf-8")
+                    else:
+                        return f.read()
-- 
2.30.2