From e7e13f13805fe833dc9852e986584b6469c5060f Mon Sep 17 00:00:00 2001
From: ale <ale@incal.net>
Date: Sat, 10 Dec 2016 12:33:52 +0000
Subject: [PATCH] Make urllib2_handler more user-friendly.

Allows the caller to specify password and other auth parameters, and to
force sessions to be non-interactive (for scripts and such).
---
 src/python/sso/urllib2_handler.py | 73 +++++++++++++++++++++----------
 1 file changed, 50 insertions(+), 23 deletions(-)

diff --git a/src/python/sso/urllib2_handler.py b/src/python/sso/urllib2_handler.py
index ee58444..4f60ee9 100644
--- a/src/python/sso/urllib2_handler.py
+++ b/src/python/sso/urllib2_handler.py
@@ -7,16 +7,36 @@ import re
 import sys
 
 
-class SSOProcessor(urllib2.BaseHandler):
+DEFAULT_LOGIN_SERVER = 'https://login.autistici.org/'
+
 
-    LOGIN_SERVER = 'https://login.autistici.org/'
+class SSOProcessor(urllib2.BaseHandler):
+    """Intercept SSO login requests and fulfills them on-the-fly."""
 
     _form_pattern = re.compile(r'<input type="hidden" name="([^"]+)" value="([^"]+)"')
     _otp_pattern = re.compile(r'<input[^>]+ name="otp"')
 
-    def __init__(self, username=None, password=None):
+    def __init__(self, username=None, password=None, otp=None, login_server=None,
+                 interactive=None):
+        """Constructor.
+
+        Args:
+          username: username, defaults to current user if unspecified.
+          password: password, if already known. Will ask interactively if
+               possible.
+          otp: OTP token, if already known. In most cases it makes no sense
+               to pass this argument given its dependency on the current time.
+          login_server: login server URL, defaults to DEFAULT_LOGIN_SERVER.
+          interactive: force the session to be considered interactive or not
+               (by default we will look at stdin.isatty).
+        """
         self._username = username or getpass.getuser()
         self._password = password
+        self._otp = otp
+        self._login_server = login_server or DEFAULT_LOGIN_SERVER
+        if interactive is None:
+            interactive = sys.stdin.isatty()
+        self._interactive = interactive
 
     def _extract_hidden_form_data(self, html):
         form = {}
@@ -29,36 +49,43 @@ class SSOProcessor(urllib2.BaseHandler):
 
     def https_response(self, req, resp):
         request_url = req.get_full_url()
-        if (resp.code == 200 and request_url.startswith(self.LOGIN_SERVER)):
-            if not hasattr(req, 'sso_attempt'):
-                request_baseurl = request_url.split('?')[0]
-                response_data = resp.read()
-                form_data = self._extract_hidden_form_data(response_data)
-                form_data['username'] = self._username
-                if not self._password:
+        if (resp.code == 200 and request_url.startswith(self._login_server)
+            and not hasattr(req, 'sso_attempt')):
+            request_baseurl = request_url.split('?')[0]
+            response_data = resp.read()
+            form_data = self._extract_hidden_form_data(response_data)
+            form_data['username'] = self._username
+            if not self._password:
+                if self._interactive:
                     self._password = getpass.getpass(
                         prompt='Password for %s@%s: ' % (
                             self._username, form_data['s'].rstrip('/')))
-                form_data['password'] = self._password
-                if self._otp_pattern.search(response_data):
-                    # Only ask for an OTP when the standard input is a tty,
-                    # otherwise simply don't send any OTP.
-                    if sys.stdin.isatty():
-                        form_data['otp'] = raw_input('OTP for %s@%s: ' % (
-                            self._username, form_data['s'].rstrip('/')))
-                newreq = urllib2.Request(request_baseurl, 
-                                         data=urllib.urlencode(form_data))
-                newreq.sso_attempt = True
-                resp = self.parent.open(newreq)
+                else:
+                    raise Exception('No password available')
+            form_data['password'] = self._password
+            # See if the form is requesting an OTP token.
+            if self._otp_pattern.search(response_data):
+                # Only ask for an OTP when the standard input is a tty,
+                # otherwise simply don't send any OTP.
+                otp = self._otp
+                if otp is None and self._interactive:
+                    otp = raw_input('OTP for %s@%s: ' % (
+                        self._username, form_data['s'].rstrip('/')))
+                if otp:
+                    form_data['otp'] = otp
+            newreq = urllib2.Request(request_baseurl,
+                                     data=urllib.urlencode(form_data))
+            newreq.sso_attempt = True
+            resp = self.parent.open(newreq)
         return resp
 
 
-def install_handler(username=None, jar=None):
+def install_handler(jar=None, **kwargs):
     if jar is None:
         jar = cookielib.CookieJar()
     urllib2.install_opener(
         urllib2.build_opener(urllib2.HTTPCookieProcessor(jar),
-                             SSOProcessor(username=username)))
+                             SSOProcessor(**kwargs)))
 
 
 if __name__ == '__main__':
-- 
GitLab