fix: fix review comment

Author: fatelei

Lib/poplib.py

@@ -13,6 +13,7 @@
 
 # Imports
 
+import binascii
 import errno
 import re
 import socket
@@ -47,6 +48,8 @@ class error_proto(Exception): pass
 # 512 characters, including CRLF. We have selected 2048 just to be on
 # the safe side.
 _MAXLINE = 2048
+# maximum number of AUTH challenges we are willing to process (parity with smtplib)
+_MAXCHALLENGE = 5
 
 
 class POP3:
@@ -218,61 +221,6 @@ def pass_(self, pswd):
         """
         return self._shortcmd('PASS %s' % pswd)
 
-    def auth(self, mechanism, authobject=None, initial_response=None):
-        """Authenticate to the POP3 server using the AUTH command (RFC 5034).
-
-        Result is 'response'.
-        """
-        if authobject is not None and initial_response is not None:
-            raise ValueError('authobject and initial_response are mutually exclusive')
-
-        if initial_response is not None:
-            if isinstance(initial_response, str):
-                initial_response = initial_response.encode(self.encoding)
-            b64 = base64.b64encode(initial_response).decode('ascii')
-            return self._shortcmd(f'AUTH {mechanism} {b64}'.rstrip())
-
-        if authobject is None:
-            return self._shortcmd(f'AUTH {mechanism}')
-
-        self._putcmd(f'AUTH {mechanism}')
-        while True:
-            line, _ = self._getline()
-            if line.startswith(b'+OK'):
-                return line
-            if line.startswith(b'-ERR'):
-                while self._getline() != b'.\r\n':
-                    pass
-                raise error_proto(line.decode('ascii', 'replace'))
-
-            if not line.startswith(b'+ '):
-                raise error_proto(f'malformed challenge line: {line!r}')
-
-            challenge_b64 = line[2:]
-            challenge_b64 = challenge_b64.rstrip(b'\r\n')
-
-            if challenge_b64:
-                try:
-                    challenge = base64.b64decode(challenge_b64)
-                except Exception:
-                    padded = challenge_b64 + b'=' * (-len(challenge_b64) % 4)
-                    challenge = base64.b64decode(padded)
-            else:
-                challenge = b''
-
-            response = authobject(challenge)
-            if response is None:
-                response = b''
-            if isinstance(response, str):
-                response = response.encode(self.encoding)
-
-            if response == b'*':
-                self._putcmd('*')
-                err_line, _ = self._getline()
-                raise error_proto(err_line.decode('ascii', 'replace'))
-
-            self._putcmd(base64.b64encode(response).decode('ascii'))
-
     def stat(self):
         """Get mailbox status.
 
@@ -479,6 +427,88 @@ def stls(self, context=None):
         self._tls_established = True
         return resp
 
+    def auth(self, mechanism, authobject, *, initial_response_ok=True):
+        """Authenticate to the POP3 server using AUTH (RFC 5034).
+
+        Result is 'response'.
+        """
+        mech = mechanism.upper()
+
+        initial = None
+        if initial_response_ok:
+            try:
+                initial = authobject()
+            except TypeError:
+                initial = None
+        if isinstance(initial, str):
+            initial = initial.encode('ascii', 'strict')
+        if initial is not None and not isinstance(initial, (bytes, bytearray)):
+            raise TypeError('authobject() must return str or bytes for initial response')
+
+        if initial is not None:
+            b64 = base64.b64encode(initial).decode('ascii')
+            cmd = f'AUTH {mech} {b64}'
+            if len(cmd.encode('ascii')) + 2 <= 255:
+                self._putcmd(cmd)
+            else:
+                self._putcmd(f'AUTH {mech}') 
+        else:
+            self._putcmd(f'AUTH {mech}')
+
+        auth_challenge_count = 0
+        while True:
+            line, _ = self._getline()
+
+            if line.startswith(b'+OK'):
+                return line
+
+            if line.startswith(b'-ERR'):
+                raise error_proto(line.decode('ascii', 'replace'))
+            # Challenge line: "+ <b64>" or just "+" (empty challenge)
+            if not (line == b'+' or line.startswith(b'+ ')):
+                raise error_proto(f'malformed AUTH challenge line: {line!r}')
+
+            auth_challenge_count += 1
+            if auth_challenge_count > _MAXCHALLENGE:
+                raise error_proto('Server AUTH mechanism infinite loop')
+
+            chal = line[1:]
+            if chal.startswith(b' '):
+                chal = chal[1:]
+            chal = chal.rstrip(b'\r\n')
+            if chal:
+                try:
+                    challenge = base64.b64decode(chal, validate=True)
+                except (binascii.Error, ValueError):
+                    self._putcmd('*')
+                    line, _ = self._getline()
+                    raise error_proto(line.decode('ascii', 'replace'))
+            else:
+                challenge = b''
+
+            resp = authobject(challenge)
+            if resp is None:
+                resp = b''
+            if isinstance(resp, str):
+                resp = resp.encode('ascii', 'strict')
+            if not isinstance(resp, (bytes, bytearray)):
+                raise TypeError('authobject(challenge) must return str or bytes')
+
+            if resp == b'*':
+                self._putcmd('*')
+            else:
+                self._putcmd(base64.b64encode(resp).decode('ascii'))
+
+    def auth_plain(self, user, password, authzid=''):
+        """Return an authobject suitable for SASL PLAIN.
+
+        Result is 'str'.
+        """
+        def _auth_plain(challenge=None):
+            # Per RFC 4616, the response is: authzid UTF8 NUL authcid UTF8 NUL passwd UTF8
+            return f"{authzid}\0{user}\0{password}"
+        return _auth_plain
+
 
 if HAVE_SSL:
 
@@ -513,6 +543,7 @@ def stls(self, context=None):
             SSL/TLS session.
             """
             raise error_proto('-ERR TLS session already established')
+        
 
     __all__.append("POP3_SSL")
 

Lib/test/test_poplib.py

@@ -326,7 +326,7 @@ def test_pass_(self):
 
     def test_auth_plain_initial_response(self):
         secret = b"user\x00adminuser\x00password"
-        resp = self.client.auth("PLAIN", initial_response=secret)
+        resp = self.client.auth("PLAIN", authobject=lambda: secret)
         self.assertStartsWith(resp, b"+OK")
 
     def test_auth_plain_challenge_response(self):
@@ -336,13 +336,9 @@ def authobject(challenge):
         resp = self.client.auth("PLAIN", authobject=authobject)
         self.assertStartsWith(resp, b"+OK")
 
-    def test_auth_rejects_conflicting_args(self):
-        with self.assertRaises(ValueError):
-            self.client.auth("PLAIN", authobject=lambda c: b"x", initial_response=b"y")
-
     def test_auth_unsupported_mechanism(self):
         with self.assertRaises(poplib.error_proto):
-            self.client.auth("FOO")
+            self.client.auth("FOO", authobject=lambda: b"")
 
     def test_auth_cancel(self):
         def authobject(_challenge):
@@ -353,12 +349,12 @@ def authobject(_challenge):
     def test_auth_mechanism_case_insensitive(self):
         secret = b"user\x00adminuser\x00password"
         # use lowercase mechanism name to ensure server accepts
-        resp = self.client.auth("plain", initial_response=secret)
+        resp = self.client.auth("plain", authobject=lambda: secret)
         self.assertStartsWith(resp, b"+OK")
 
     def test_auth_initial_response_str(self):
         secret = "user\x00adminuser\x00password"  # str, not bytes
-        resp = self.client.auth("PLAIN", initial_response=secret)
+        resp = self.client.auth("PLAIN", authobject=lambda: secret)
         self.assertStartsWith(resp, b"+OK")
 
     def test_auth_authobject_returns_str(self):