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

@@ -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,23 +49,30 @@ 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'):
+        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('/')))
+                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.
-                    if sys.stdin.isatty():
-                        form_data['otp'] = raw_input('OTP for %s@%s: ' % (
+                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
@@ -53,12 +80,12 @@ class SSOProcessor(urllib2.BaseHandler):
         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__':