From 7fed8de8f04d263cc7f704d5d16533cc8676a350 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Matigot?= <fred@fspot.org>
Date: Mon, 22 Feb 2021 09:54:05 +0100
Subject: [PATCH] feat: handle weird ftp server (#40)

* feat: handle weird ftp server

* chore: meh

* chore: PROTOCOL_TLSv1 is not supported on ubuntu 20.04

* chore: fix for mypy

* chore: ignore line from coverage

* chore: bump pypi version number
---
 .github/workflows/ci.yml    |  4 ++--
 peakina/io/ftp/ftp_utils.py | 28 ++++++++++++++++++++++++----
 setup.cfg                   |  2 +-
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 6a8bc7be..86c08f6f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -7,13 +7,13 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8]
+        python-version: ['3.6', '3.7', '3.8']
 
     steps:
     - uses: actions/checkout@v2
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
 
diff --git a/peakina/io/ftp/ftp_utils.py b/peakina/io/ftp/ftp_utils.py
index ff17d823..d41a3a0c 100644
--- a/peakina/io/ftp/ftp_utils.py
+++ b/peakina/io/ftp/ftp_utils.py
@@ -7,6 +7,7 @@
 from contextlib import contextmanager, suppress
 from datetime import datetime
 from functools import partial
+from ipaddress import ip_address
 from os.path import basename, join
 from time import sleep
 from typing import Dict, List, Optional
@@ -19,6 +20,8 @@
 
 
 class FTPS(ftplib.FTP_TLS):
+    ssl_version = ssl.PROTOCOL_TLSv1_2
+
     def connect(self, host, port, timeout):
         self.host = host
         self.port = port or 990
@@ -26,13 +29,30 @@ def connect(self, host, port, timeout):
 
         self.sock = socket.create_connection((self.host, self.port), self.timeout)
         self.af = self.sock.family
-        self.sock = ssl.wrap_socket(
-            self.sock, self.keyfile, self.certfile, ssl_version=ssl.PROTOCOL_TLSv1
-        )
+        self.sock = self.context.wrap_socket(self.sock, server_hostname=self.host)
         self.file = self.sock.makefile('r')
         self.welcome = self.getresp()
         return self.welcome
 
+    def ntransfercmd(self, cmd, rest=None):
+        # override ntransfercmd so it reuses the sock session, to prevent SSLEOFError.
+        # cf. https://stackoverflow.com/questions/40536061/ssleoferror-on-ftps-using-python-ftplib
+        conn, size = ftplib.FTP.ntransfercmd(self, cmd, rest)
+        if self._prot_p:
+            conn = self.context.wrap_socket(
+                conn, server_hostname=self.host, session=self.sock.session
+            )  # this is the fix
+        return conn, size
+
+    def makepasv(self):
+        # override makepasv so it rewrites the dst address if the server gave a broken one.
+        # Inspired by:
+        # https://github.com/lavv17/lftp/blob/d67fc14d085849a6b0418bb3e912fea2e94c18d1/src/ftpclass.cc#L774
+        host, port = super().makepasv()
+        if self.af == socket.AF_INET and ip_address(host).is_private:  # pragma: no cover
+            host = self.sock.getpeername()[0]
+        return host, port
+
 
 @contextmanager
 def ftps_client(params: ParseResult):
@@ -69,7 +89,7 @@ def sftp_client(params: ParseResult):
     try:
         ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
         ssh_client.connect(
-            hostname=params.hostname,
+            hostname=params.hostname,  # type: ignore
             username=params.username,
             password=params.password,
             port=port,
diff --git a/setup.cfg b/setup.cfg
index 459883a1..e78cde2c 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -6,7 +6,7 @@ long-description-content-type = text/markdown; charset=UTF-8
 author = Toucan Toco
 author_email = dev@toucantoco.com
 url = https://github.com/ToucanToco/peakina
-version = 0.5.5
+version = 0.5.6
 license = BSD
 classifiers=
     Intended Audience :: Developers