Commit e7e13f13 authored by ale's avatar ale

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).
parent 3078e5c6
...@@ -7,16 +7,36 @@ import re ...@@ -7,16 +7,36 @@ import re
import sys import sys
class SSOProcessor(urllib2.BaseHandler): DEFAULT_LOGIN_SERVER = ''
LOGIN_SERVER = '' class SSOProcessor(urllib2.BaseHandler):
"""Intercept SSO login requests and fulfills them on-the-fly."""
_form_pattern = re.compile(r'<input type="hidden" name="([^"]+)" value="([^"]+)"') _form_pattern = re.compile(r'<input type="hidden" name="([^"]+)" value="([^"]+)"')
_otp_pattern = re.compile(r'<input[^>]+ name="otp"') _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,
username: username, defaults to current user if unspecified.
password: password, if already known. Will ask interactively if
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._username = username or getpass.getuser()
self._password = password 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): def _extract_hidden_form_data(self, html):
form = {} form = {}
...@@ -29,36 +49,43 @@ class SSOProcessor(urllib2.BaseHandler): ...@@ -29,36 +49,43 @@ class SSOProcessor(urllib2.BaseHandler):
def https_response(self, req, resp): def https_response(self, req, resp):
request_url = req.get_full_url() request_url = req.get_full_url()
if (resp.code == 200 and request_url.startswith(self.LOGIN_SERVER)): if (resp.code == 200 and request_url.startswith(self._login_server)
if not hasattr(req, 'sso_attempt'): and not hasattr(req, 'sso_attempt')):
request_baseurl = request_url.split('?')[0] request_baseurl = request_url.split('?')[0]
response_data = response_data =
form_data = self._extract_hidden_form_data(response_data) form_data = self._extract_hidden_form_data(response_data)
form_data['username'] = self._username form_data['username'] = self._username
if not self._password: if not self._password:
if self._interactive:
self._password = getpass.getpass( self._password = getpass.getpass(
prompt='Password for %s@%s: ' % ( prompt='Password for %s@%s: ' % (
self._username, form_data['s'].rstrip('/'))) self._username, form_data['s'].rstrip('/')))
form_data['password'] = self._password else:
if raise Exception('No password available')
# Only ask for an OTP when the standard input is a tty, form_data['password'] = self._password
# otherwise simply don't send any OTP. # See if the form is requesting an OTP token.
if sys.stdin.isatty(): if
form_data['otp'] = raw_input('OTP for %s@%s: ' % ( # Only ask for an OTP when the standard input is a tty,
self._username, form_data['s'].rstrip('/'))) # otherwise simply don't send any OTP.
newreq = urllib2.Request(request_baseurl, otp = self._otp
data=urllib.urlencode(form_data)) if otp is None and self._interactive:
newreq.sso_attempt = True otp = raw_input('OTP for %s@%s: ' % (
resp = self._username, form_data['s'].rstrip('/')))
if otp:
form_data['otp'] = otp
newreq = urllib2.Request(request_baseurl,
newreq.sso_attempt = True
resp =
return resp return resp
def install_handler(username=None, jar=None): def install_handler(jar=None, **kwargs):
if jar is None: if jar is None:
jar = cookielib.CookieJar() jar = cookielib.CookieJar()
urllib2.install_opener( urllib2.install_opener(
urllib2.build_opener(urllib2.HTTPCookieProcessor(jar), urllib2.build_opener(urllib2.HTTPCookieProcessor(jar),
SSOProcessor(username=username))) SSOProcessor(**kwargs)))
if __name__ == '__main__': if __name__ == '__main__':
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment