From 1da8002321cccd317349cdf9b888876832c69a19 Mon Sep 17 00:00:00 2001
From: Oleg Broytman <phd@phdru.name>
Date: Tue, 6 Aug 2024 14:27:46 +0300
Subject: [PATCH] Feat(bkmk_raiohttp): Use aioftp

---
 Robots/bkmk_raiohttp.py | 42 +++++++++++++++++++++++++++++++++++++++--
 Robots/bkmk_rcurl.py    |  2 +-
 bkmk_db-venv            |  3 ++-
 doc/ANNOUNCE            |  2 ++
 doc/ChangeLog           |  2 ++
 doc/TODO                |  2 --
 setup.py                |  2 +-
 7 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/Robots/bkmk_raiohttp.py b/Robots/bkmk_raiohttp.py
index ed4dac6..e94d72a 100644
--- a/Robots/bkmk_raiohttp.py
+++ b/Robots/bkmk_raiohttp.py
@@ -11,9 +11,13 @@ __license__ = "GNU GPL"
 __all__ = ['robot_aiohttp']
 
 
+from urllib.parse import urlsplit
 import asyncio
+
+import aioftp
 import aiohttp
 import aiohttp.client_exceptions
+
 from Robots.bkmk_robot_base import robot_base, request_headers
 
 
@@ -22,6 +26,16 @@ class robot_aiohttp(robot_base):
         return 'aiohttp/%s' % aiohttp.__version__
 
     def get(self, bookmark, url, accept_charset=False, use_proxy=False):
+        if url.startswith('ftp://'):
+            error, _, _, body = asyncio.run(get_ftp(
+                url, connect_timeout=self.connect_timeout,
+                timeout=self.timeout,
+            ))
+            if error is not None:
+                error = str(error)
+                return error, None, None, None, None
+            return None, None, None, None, body
+
         if accept_charset and bookmark.charset:
             headers = request_headers.copy()
             headers['Accept-Charset'] = bookmark.charset
@@ -33,7 +47,7 @@ class robot_aiohttp(robot_base):
         else:
             proxy = None
 
-        error, status, resp_headers, body = asyncio.run(get(
+        error, status, resp_headers, body = asyncio.run(get_http(
             url, headers=headers, proxy=proxy,
             connect_timeout=self.connect_timeout, timeout=self.timeout,
         ))
@@ -49,8 +63,11 @@ class robot_aiohttp(robot_base):
             return None, status, resp_headers['Location'], None, None
         return None, status, None, resp_headers, body
 
+    def get_ftp_welcome(self):
+        return ''  # We don't store welcome message yet
+
 
-async def get(url, headers={}, proxy=None, connect_timeout=30, timeout=60):
+async def get_http(url, headers={}, proxy=None, connect_timeout=30, timeout=60):
     timeout = aiohttp.ClientTimeout(connect=connect_timeout, total=timeout)
     try:
         async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -60,3 +77,24 @@ async def get(url, headers={}, proxy=None, connect_timeout=30, timeout=60):
                 return None, resp.status, resp.headers, await resp.read()
     except (asyncio.TimeoutError, aiohttp.client_exceptions.ClientError) as e:
         return e, None, None, None
+
+
+async def get_ftp(url, connect_timeout=30, timeout=60):
+    split_results = urlsplit(url)
+    path = split_results.path or '/'
+    user = split_results.username or 'anonymous'
+    password = split_results.password or 'ftp'
+    host = split_results.hostname
+    port = split_results.port or 21
+
+    lines = []
+    try:
+        async with aioftp.Client.context(
+            host, port, user=user, password=password,
+            socket_timeout=connect_timeout, path_timeout=timeout,
+        ) as client:
+            async for _path, _info in client.list(path, recursive=True):
+                lines.append('%s %s' % (_info, _path))
+    except (OSError, aioftp.errors.AIOFTPException) as e:
+        return e, None, None, None
+    return None, None, None, '\n'.join(lines)
diff --git a/Robots/bkmk_rcurl.py b/Robots/bkmk_rcurl.py
index bbbc9d1..bf32970 100644
--- a/Robots/bkmk_rcurl.py
+++ b/Robots/bkmk_rcurl.py
@@ -95,7 +95,7 @@ class robot_curl(robot_base):
         self.body += data
 
     def get_ftp_welcome(self):
-        return ''  # We doen't store welcome message yet
+        return ''  # We don't store welcome message yet
 
 
 def encode_url(url, encoding):
diff --git a/bkmk_db-venv b/bkmk_db-venv
index 62d0d98..6b8f65e 100644
--- a/bkmk_db-venv
+++ b/bkmk_db-venv
@@ -9,6 +9,7 @@ if [ -z "$VIRTUAL_ENV" ]; then
          . bkmk_db-venv/bin/activate &&
          pip install --compile --upgrade beautifulsoup4 lxml m_lib.full \
          requests requests-ftp \
-         pycurl certifi aiohttp
+         pycurl certifi \
+         aiohttp aioftp
     }
 fi
diff --git a/doc/ANNOUNCE b/doc/ANNOUNCE
index ec435e1..bb7b34a 100644
--- a/doc/ANNOUNCE
+++ b/doc/ANNOUNCE
@@ -9,6 +9,8 @@ WHAT'S NEW
 
 Version 5.5.1 (2024-08-??)
 
+   Use aioftp in aiohttp robot.
+
    Do not route ftp requests via http(s) proxy; socks5 proxies are ok.
 
 Version 5.5.0 (2024-08-06)
diff --git a/doc/ChangeLog b/doc/ChangeLog
index 4057f50..0d40430 100644
--- a/doc/ChangeLog
+++ b/doc/ChangeLog
@@ -1,5 +1,7 @@
 Version 5.5.1 (2024-08-??)
 
+   Use aioftp in aiohttp robot.
+
    Do not route ftp requests via http(s) proxy; socks5 proxies are ok.
 
 Version 5.5.0 (2024-08-06)
diff --git a/doc/TODO b/doc/TODO
index d2e6573..ad7ddd2 100644
--- a/doc/TODO
+++ b/doc/TODO
@@ -1,5 +1,3 @@
-aioftp.
-
 Robot(s) that test many URLs in parallel.
 
 A program to publish bookmarks with icons.
diff --git a/setup.py b/setup.py
index 0a570ad..f6f6e01 100755
--- a/setup.py
+++ b/setup.py
@@ -41,6 +41,6 @@ setup(
         'html': ['beautifulsoup4', 'lxml'],
         'requests': ['requests', 'requests-ftp'],
         'curl': ['pycurl', 'certifi'],
-        'aiohttp:python_version>="3.4"': ['aiohttp'],
+        'aiohttp:python_version>="3.4"': ['aiohttp', 'aioftp'],
     },
 )
-- 
2.39.5