Merge pull request #443 from OP-TED/feature/TED-1247

add private key auth method support for SFTP

Author: CaptainOfHacks

requirements.txt

@@ -20,10 +20,9 @@ agraph-python==101.0.10
 decorator~=5.1.1
 urllib3[secure]
 semantic-version==2.10.0
-pysftp~=0.2.9
+paramiko~=3.0.0
 ordered-set~=4.0.2
 json2html~=1.3.0
 minio~=7.1.1
-certifi
-networkx~=2.8.8
+certifi~=2022.12.7
 shortuuid~=1.0.11

ted_sws/__init__.py

@@ -12,7 +12,7 @@
 import json
 import os
 import pathlib
-
+import base64
 import dotenv
 
 from ted_sws.core.adapters.config_resolver import EnvConfigResolver, AirflowAndEnvConfigResolver, env_property
@@ -201,6 +201,21 @@ def SFTP_PUBLISH_PASSWORD(self, config_value: str) -> str:
     def SFTP_PUBLISH_PATH(self, config_value: str) -> str:
         return config_value
 
+    @env_property(config_resolver_class=AirflowAndEnvConfigResolver)
+    def SFTP_PRIVATE_KEY_BASE64(self, config_value: str) -> str:
+        return config_value
+
+    @env_property(config_resolver_class=AirflowAndEnvConfigResolver)
+    def SFTP_PRIVATE_KEY_PASSPHRASE(self, config_value: str) -> str:
+        return config_value
+
+    @property
+    def SFTP_PRIVATE_KEY(self) -> str:
+        sftp_private_key_base64 = self.SFTP_PRIVATE_KEY_BASE64
+        if sftp_private_key_base64:
+            sftp_private_key_base64 = base64.b64decode(str(sftp_private_key_base64)).decode(encoding="utf-8")
+        return sftp_private_key_base64
+
 
 class SPARQLConfig:
 

ted_sws/notice_publisher/adapters/sftp_notice_publisher.py

@@ -1,5 +1,6 @@
-import pysftp
+import io
 
+import paramiko
 from ted_sws import config
 from ted_sws.notice_publisher.adapters.sftp_publisher_abc import SFTPPublisherABC
 
@@ -9,36 +10,43 @@ class SFTPPublisher(SFTPPublisherABC):
 
     """
 
-    def __init__(self, hostname: str = None, username: str = None, password: str = None, port: int = None):
+    def __init__(self, hostname: str = None, username: str = None, password: str = None, port: int = None,
+                 private_key: str = None, private_key_passphrase: str = None):
         """Constructor Method"""
         self.hostname = hostname if hostname else config.SFTP_PUBLISH_HOST
         self.username = username if username else config.SFTP_PUBLISH_USER
         self.password = password if password else config.SFTP_PUBLISH_PASSWORD
         self.port = port if port else config.SFTP_PUBLISH_PORT
         self.connection = None
         self.is_connected = False
+        self.private_key = None
+        self.private_key_passphrase = private_key_passphrase if private_key_passphrase else config.SFTP_PRIVATE_KEY_PASSPHRASE
+        private_key = private_key if private_key else config.SFTP_PRIVATE_KEY
+        if private_key:
+            self.private_key = paramiko.RSAKey.from_private_key(io.StringIO(private_key),
+                                                                password=self.private_key_passphrase)
+
+        self._sftp = None
+        self._ssh = None
 
     def connect(self):
         """Connects to the sftp server and returns the sftp connection object"""
 
-        cnopts = pysftp.CnOpts()
-        # TODO: to be checked/removed when SSL will be setup
-        cnopts.hostkeys = None
-        # Get the sftp connection object
-        self.connection = pysftp.Connection(
-            host=self.hostname,
-            username=self.username,
-            password=self.password,
-            port=self.port,
-            cnopts=cnopts
-        )
-
+        ssh_client = paramiko.SSHClient()
+        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh_client.connect(self.hostname, username=self.username,
+                           pkey=self.private_key, port=self.port)
+        self._ssh = ssh_client
+        self._sftp = ssh_client.open_sftp()
         self.is_connected = True
 
     def disconnect(self):
         """Closes the sftp connection"""
         if self.is_connected:
-            self.connection.close()
+            self._sftp.close()
+            self._ssh.close()
+            self._sftp = None
+            self._ssh = None
             self.is_connected = False
 
     def __del__(self):
@@ -48,12 +56,25 @@ def publish(self, source_path: str, remote_path: str) -> bool:
         """
         Publish file_content to the sftp server remote path.
        """
-        self.connection.put(source_path, remote_path)
+        self._sftp.put(source_path, remote_path)
         return True
 
     def remove(self, remote_path: str) -> bool:
         """
 
         """
-        self.connection.unlink(remote_path)
+        self._sftp.remove(remote_path)
         return True
+
+    def exists(self, remote_path: str) -> bool:
+        """
+            Check if remote_path exists.
+        :param remote_path:
+        :return:
+        """
+        if self.is_connected:
+            try:
+                self._sftp.stat(remote_path)
+            except IOError:
+                return False
+        return True

tests/e2e/notice_publisher/adapters/test_notice_publisher.py

@@ -32,11 +32,12 @@ def test_sftp_notice_publisher():
     with pytest.raises(Exception):
         sftp_publisher.publish(source_file.name, None)
 
+    assert not sftp_publisher.exists(remote_path)
     published = sftp_publisher.publish(source_file.name, remote_path)
     assert published
-    assert sftp_publisher.connection.exists(remote_path)
+    assert sftp_publisher.exists(remote_path)
     sftp_publisher.remove(remote_path)
-    assert not sftp_publisher.connection.exists(remote_path)
+    assert not sftp_publisher.exists(remote_path)
 
     sftp_publisher.disconnect()
 

tests/features/notice_publisher/test_notice_publisher.py

@@ -81,7 +81,7 @@ def the_mets_package_available_in_a_shared_sftp_drive(published_notice: Notice,
     publisher: SFTPPublisher = SFTPPublisher()
     remote_notice_path = f"{sftp_remote_folder_path}/{published_notice.ted_id}.zip"
     publisher.connect()
-    assert publisher.connection.exists(remotepath=remote_notice_path)
+    assert publisher.exists(remote_path=remote_notice_path)
     publisher.disconnect()