fix FTP connectin error

enerBit · Sep 4, 2024 · b909d4d · b909d4d
1 parent 8dad523
commit b909d4d
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 7 deletions.
diff --git a/src/asic/cli.py b/src/asic/cli.py
@@ -72,6 +72,7 @@ def main(
     ctx.meta["ASIC_FTPS_PORT"] = ftps_port
     ctx.meta["ASIC_FTPS_USER"] = ftps_user
     ctx.meta["ASIC_FTPS_PASSWORD"] = pydantic.SecretStr(ftps_password)
+    ctx.meta["VERBOSITY"] = verbosity
 
 
 def validate_month(month: str) -> str:
@@ -207,6 +208,7 @@ def list_files(
     ftps_port = ctx.meta["ASIC_FTPS_PORT"]
     ftps_user = ctx.meta["ASIC_FTPS_USER"]
     ftps_password = ctx.meta["ASIC_FTPS_PASSWORD"]
+    verbosity = ctx.meta["VERBOSITY"]
 
     if not extensions:
         extensions = [None]  # type: ignore
@@ -222,6 +224,7 @@ def list_files(
         ftps_user=ftps_user,
         ftps_password=ftps_password,
         ftps_port=ftps_port,
+        verbosity=verbosity,
     )
     month_dates = [parse_month(m) for m in months]
     file_list = list_supported_files(
@@ -282,6 +285,7 @@ def download(
     ftps_port = ctx.meta["ASIC_FTPS_PORT"]
     ftps_user = ctx.meta["ASIC_FTPS_USER"]
     ftps_password = ctx.meta["ASIC_FTPS_PASSWORD"]
+    verbosity = ctx.meta["VERBOSITY"]
 
     locations = PUBLIC_SEARCHEABLE_LOCATIONS
     if agent is not None:
@@ -292,6 +296,7 @@ def download(
         ftps_user=ftps_user,
         ftps_password=ftps_password,
         ftps_port=ftps_port,
+        verbosity=verbosity,
     )
     month_dates = [parse_month(m) for m in months]
     file_list = list_supported_files(
@@ -323,6 +328,7 @@ def download(
                 ftps_user=ftps_user,
                 ftps_password=ftps_password,
                 ftps_port=ftps_port,
+                verbosity=verbosity,
             )
 
             grab_file(ftps, remote.path, local)

diff --git a/src/asic/ftp.py b/src/asic/ftp.py
@@ -1,9 +1,12 @@
+import contextlib
 import datetime as dt
 import ftplib
 import functools
+import io
 import itertools
 import logging
 import pathlib
+import ssl
 from typing import Iterable, Type
 
 import pydantic
@@ -14,6 +17,9 @@
 
 logger = logging.getLogger(__name__)
 
+
+SSL_CONTEXT = ssl.create_default_context()
+
 SUPPORTED_ASIC_EXTENSIONS = frozenset(ASIC_FILE_EXTENSION_MAP.keys())
 
 
@@ -30,21 +36,34 @@ def get_ftps(
     ftps_user: str,
     ftps_password: pydantic.SecretStr,
     ftps_port: int,
+    verbosity: int = 0,
 ):
     ftps = ftplib.FTP_TLS(
-        host=ftps_host,
         encoding="Latin-1",
-        # timeout=10,
     )
 
-    ftps.connect(
-        port=ftps_port,
-    )
+    ftps.set_debuglevel(verbosity - 1)
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        ftps.connect(
+            host=ftps_host,
+            port=ftps_port,
+        )
+    for line in f.getvalue().splitlines():
+        logger.debug(f"FTP: {line}")
     logger.info(f"Login to FTP '{ftps_host}' as '{ftps_user}'")
-    ftps.login(user=ftps_user, passwd=ftps_password.get_secret_value())
 
-    ftps.prot_p()
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        ftps.login(user=ftps_user, passwd=ftps_password.get_secret_value())
+    for line in f.getvalue().splitlines():
+        logger.debug(f"FTP: {line}")
+
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        # Explicitly secure the connection
+        ftps.prot_p()
 
+    for line in f.getvalue().splitlines():
+        logger.debug(f"FTP: {line}")
+    ftps.set_debuglevel(0)
     return ftps