box-oauth-tfa.py

#!/usr/local/bin/python

import urllib2
import urllib
import cookielib
from ntlm import HTTPNtlmAuthHandler
import StringIO
import lxml.html
import json
import getpass

user = raw_input('Username: ')
pwd = getpass.getpass('Password: ')
domain = raw_input('Domain: ')
box_user = raw_input('Box user: ')

client_id = ''
client_secret = ''

authurl = 'https://app.box.com'
authuri = '/api/oauth2/authorize?response_type=code&client_id=' + client_id
ssourl = 'https://sso.services.box.net'
twofactorurl = 'https://...'
ntlmrealm = 'https://...'

def doAUTH(url, opener):
    first = HTTPrequest(opener, 
        url)
    firsthtml = parseHTMLelements(first)
    formdict = httpRedirectURL(firsthtml)
    second = HTTPrequest(opener, 
        authurl + formdict['_redirect_url'], 
        params=formdict)
    url, relaystate, samlrequest = httpSAMLrequest(parseHTMLelements(second))
    third = HTTPrequest(opener, 
        url, 
        params={
            "SAMLRequest":samlrequest,
            "RelayState":relaystate
            }
        )
    thirdhtml = parseHTMLelements(third)
    pinform = pinSubmit(thirdhtml)
    forth = HTTPrequest(opener, 
        twofactorurl + "/" + thirdhtml.xpath("//form[@method='post']")[0].action, 
        params=pinform)
    forthhtml = parseHTMLelements(forth)
    url, relaystate, samlresponse = httpSAMLresponse(forthhtml)
    fifth = HTTPrequest(opener, 
        url, 
        params={
            "SAMLResponse":samlresponse,
            "RelayState":relaystate
            }
        )
    fifthhtml = parseHTMLelements(fifth)
    tokenform = opentoken(fifthhtml)
    sixth = HTTPrequest(opener, 
        fifthhtml.xpath("//form[@method='post']")[0].action, 
        params=tokenform)
    sixthhtml = parseHTMLelements(sixth)
    consentform = consentForm(sixthhtml)
    seventh = HTTPrequest(rediropener, 
        sixthhtml.xpath("//form[@method='post']")[0].action, 
        params=consentform)
    oauthcode = seventh['location'].split("=")[2]
    eighth = HTTPrequest(opener, 
        authurl + "/api/oauth2/token", 
        params={
            "grant_type":"authorization_code", 
            "code":oauthcode, 
            "client_id":client_id, 
            "client_secret":client_secret
            }
        )
    jsonout = json.loads(eighth.read()) # response contains a JSON object with access_token, expires_in, refresh_token, token_type, restricted_to
    return jsonout

# response page-specific functions
def consentForm(consentpage):
    formdict = {}
    for i in consentpage.xpath("//input"):
        formdict[str(i.name)] = str(i.value)
    formdict['consent_accept'] = 'Grant access to Box'
    del formdict['consent_reject']
    del formdict['None']
    formdict['request_token'] = consentpage.xpath("//script")[0].text_content().split("\'")[1]
    return formdict

def opentoken(tokenpage):
    formdict = {}
    for i in tokenpage.xpath("//input"):
        formdict[str(i.name)] = str(i.value)
    return formdict

def pinSubmit(pinpage):
    formdict = {}
    pin = getpass.getpass('Enter PIN: ')
    for i in pinpage.xpath("//input"):
        formdict[str(i.name)] = str(i.value)
    formdict['ctl00$LeftSection$txtUserName'] = user
    formdict['ctl00$LeftSection$ddlDomain'] = domain
    formdict['ctl00$LeftSection$txtPassword'] = pwd
    formdict['ctl00$LeftSection$txtPasscode'] = pin
    return formdict

def httpRedirectURL(html):
    formdict = {}
    for i in html.xpath("//input"):
        formdict[str(i.name)] = str(i.value)
    formdict['login'] = str(box_user)
    formdict['login_submit'] = "Authorizing..."
    return formdict

# More generic HTTP interactions
def HTTPrequest(opener, url, params=None):
    if params is None:
        # thats a GET
        req = urllib2.Request(url)
        res = opener.open(req)
        return res
    else:
        # thats a POST
        req = urllib2.Request(url, urllib.urlencode(params))
        res = opener.open(req)
        return res

def parseHTMLelements(res):
    parser = lxml.etree.HTMLParser()
    tree = lxml.etree.parse(StringIO.StringIO(res.read()), parser)
    result = lxml.etree.tostring(tree.getroot(), pretty_print=False, method="html")
    html = lxml.html.fromstring(result)
    return html

def httpSAMLrequest(html):
    samlrequest = html.xpath("//input[@name='SAMLRequest']")[0].value
    relaystate = html.xpath("//input[@name='RelayState']")[0].value
    post = html.xpath("//form[@method='post']")[0].action
    return (post, relaystate, samlrequest)

def httpSAMLresponse(html):
    samlresponse = html.xpath("//input[@name='SAMLResponse']")[0].value
    relaystate = html.xpath("//input[@name='RelayState']")[0].value
    post = html.xpath("//form[@method='post']")[0].action
    return (post, relaystate, samlresponse)

# Overload the RedirectHandler to catch the Location: header before redirecting
class OauthHTTPRedirectHandler(urllib2.HTTPRedirectHandler):
    def http_error_302(self, req, fp, code, msg, headers):
        return headers

# NTLM handler setup
passman = urllib2.HTTPPasswordMgrWithDefaultRealm()
passman.add_password(None, ntlmrealm, user, pwd)
authNTLM = HTTPNtlmAuthHandler.HTTPNtlmAuthHandler(passman)
cj = cookielib.LWPCookieJar()
# need to build a cookie to catch the NTLM challenges with. Adjust to be company-specific
# (version, name, value, port, port_specified, domain, domain_specified, domain_initial_dot, path, path_specified, secure, expires, discard, comment, comment_url, rest, rfc2109=False)
cj.set_cookie(cookielib.Cookie(None, 'name', 'value', None, None, 'sso.company.com', None, None, '/', None, None, None, None, None, None, None))

cookie = urllib2.HTTPCookieProcessor(cj)
opener = urllib2.build_opener(cookie, urllib2.HTTPSHandler(debuglevel=0), authNTLM)
rediropener = urllib2.build_opener(cookie, urllib2.HTTPSHandler(debuglevel=1), OauthHTTPRedirectHandler)