Merge ef3d806 into 30a0b08

requirements.txt

@@ -16,3 +16,4 @@ aiodocker
 tornado
 mako
 pyjwt
+aioftp

tanner/emulators/rfi.py

@@ -1,13 +1,13 @@
 import asyncio
-import ftplib
+import aioftp
 import hashlib
 import logging
 import os
 import re
 import ssl
 import time
 from concurrent.futures import ThreadPoolExecutor
-
+from aioftp.errors import AIOFTPException, StatusCodeError, PathIsNotAbsolute, PathIOError, NoAvailablePort
 import aiohttp
 import yarl
 
@@ -26,7 +26,6 @@ def __init__(self, root_dir, loop=None, allow_insecure=False):
     async def download_file(self, path):
         file_name = None
         url = re.match(patterns.REMOTE_FILE_URL, path)
-
         if url is None:
             return None
         url = url.group(1)
@@ -37,8 +36,7 @@ async def download_file(self, path):
 
         if url.scheme == "ftp":
             pool = ThreadPoolExecutor()
-            ftp_future = self._loop.run_in_executor(pool, self.download_file_ftp, url)
-            file_name = await ftp_future
+            file_name = await self.download_file_ftp(url)
 
         else:
             ssl_context = False if self.allow_insecure else ssl.create_default_context()
@@ -56,20 +54,20 @@ async def download_file(self, path):
                     rfile.write(data.encode('utf-8'))
         return file_name
 
-    def download_file_ftp(self, url):
+    async def download_file_ftp(self, url):
         host = url.host
-        ftp_path = url.path.rsplit('/', 1)[0][1:]
+        ftp_path = url.path
         name = url.name
         try:
-            ftp = ftplib.FTP(host)
-            ftp.login()
-            ftp.cwd(ftp_path)
+            ftp = aioftp.Client()
+            await ftp.connect(host)
+            await ftp.login()
             tmp_filename = name + str(time.time())
             file_name = hashlib.md5(tmp_filename.encode('utf-8')).hexdigest()
-            with open(os.path.join(self.script_dir, file_name), 'wb') as ftp_script:
-                self.logger.debug('Saving the FTP file as %s', os.path.join(self.script_dir, file_name))
-                ftp.retrbinary('RETR %s' % name, ftp_script.write)
-        except ftplib.all_errors as ftp_errors:
+            self.logger.debug('Saving the FTP file as %s', os.path.join(self.script_dir, file_name))
+            await ftp.download(ftp_path, os.path.join(self.script_dir, file_name), write_into=True)
+            await ftp.quit()
+        except (AIOFTPException, StatusCodeError, PathIsNotAbsolute, PathIOError, NoAvailablePort) as ftp_errors:
             self.logger.exception("Problem with ftp download %s", ftp_errors)
             return None
         else:

tanner/tests/test_rfi_emulation.py

@@ -22,10 +22,8 @@ def test_http_download_fail(self):
         self.assertIsNone(filename)
 
     def test_ftp_download(self):
-        self.handler.download_file_ftp = mock.MagicMock()
         path = 'ftp://mirror.yandex.ru/archlinux/lastupdate'
-        data = self.loop.run_until_complete(self.handler.download_file(path))
-        self.handler.download_file_ftp.assert_called_with(yarl.URL(path))
+        result = self.loop.run_until_complete(self.handler.download_file_ftp(yarl.URL(path)))
 
     def test_ftp_download_fail(self):
         path = 'ftp://mirror.yandex.ru/archlinux/foobar'