Merge pull request #16 from cevoaustralia/issue/12

Fix for issue #12
cevoaustralia · Jul 24, 2017 · fd28ba3 · fd28ba3
2 parents 385694a + 0503ef7
commit fd28ba3
Show file tree

Hide file tree

Showing 4 changed files with 220 additions and 8 deletions.
diff --git a/aws_google_auth/__init__.py b/aws_google_auth/__init__.py
@@ -167,7 +167,10 @@ def handle_sms(self, sess):
         response_page = BeautifulSoup(sess.text, 'html.parser')
         challenge_url = sess.url.split("?")[0]
 
-        sms_token  = raw_input("Enter SMS token: G-") or None
+        try:
+            sms_token  = raw_input("Enter SMS token: G-") or None
+        except NameError:
+            sms_token = input("Enter SMS token: G-") or None
 
         payload = {
             'challengeId': response_page.find('input', {'name': 'challengeId'}).get('value'),
@@ -234,7 +237,10 @@ def handle_totp(self, sess):
         challenge_url = sess.url.split("?")[0]
         challenge_id = challenge_url.split("totp/")[1]
 
-        mfa_token  = raw_input("MFA token: ") or None
+        try:
+            mfa_token  = raw_input("MFA token: ") or None
+        except NameError:
+            mfa_token  = input("MFA token: ") or None
 
         if not mfa_token:
             raise ValueError("MFA token required for % but none supplied" % self.username)
@@ -263,13 +269,30 @@ def pick_one(roles):
     while True:
         for i, role in enumerate(roles):
             print("[{:>3d}] {}".format(i+1, role))
-        choice = raw_input("Type the number (1 - {:d}) of the role to assume: ".format(len(roles)))
+
+        prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(roles))
+        try:
+            choice = raw_input(prompt)
+        except NameError:
+            choice = input(prompt)
+
         try:
             num = int(choice)
-            return roles.items()[num - 1]
+            return list(roles.items())[num - 1]
         except:
             print("Invalid choice, try again")
 
+def parse_roles(doc):
+    roles = {}
+    for x in doc.xpath('//*[@Name = "https://aws.amazon.com/SAML/Attributes/Role"]//text()'):
+        if "arn:aws:iam:" not in x:
+            continue
+
+        res = x.split(',')
+        roles[res[0]] = res[1]
+
+    return roles
+
 def cli():
     parser = argparse.ArgumentParser(
         prog="aws-google-auth",
@@ -306,15 +329,24 @@ def cli():
     )
 
     if config.google_username is None:
-        config.google_username = raw_input("Google username: ")
+        try:
+            config.google_username = raw_input("Google username: ")
+        except NameError:
+            config.google_username = raw_input("Google username: ")
     else:
         print("Google username: " + config.google_username)
 
     if config.google_idp_id is None:
-        config.google_idp_id = raw_input("Google idp: ")
+        try:
+            config.google_idp_id = raw_input("Google idp: ")
+        except NameError:
+            config.google_idp_id = input("Google idp: ")
 
     if config.google_sp_id is None:
-        config.google_sp_id = raw_input("Google sp: ")
+        try:
+            config.google_sp_id = raw_input("Google sp: ")
+        except NameError:
+            config.google_sp_id = input("Google sp: ")
 
     passwd = getpass.getpass()
 
@@ -330,7 +362,7 @@ def cli():
 
     # Parse out the roles from the SAML so we can offer them as a choice
     doc = etree.fromstring(base64.b64decode(encoded_saml))
-    roles = dict([x.split(',') for x in doc.xpath('//*[@Name = "https://aws.amazon.com/SAML/Attributes/Role"]//text()')])
+    roles = parse_roles(doc)
 
     if not config.role_arn in roles:
         config.role_arn, config.provider = pick_one(roles)

diff --git a/aws_google_auth/tests/test_issue_12_extra_comma.py b/aws_google_auth/tests/test_issue_12_extra_comma.py
@@ -0,0 +1,22 @@
+import aws_google_auth
+from os import path
+from lxml import etree
+
+here = path.abspath(path.dirname(__file__))
+
+with open(path.join(here, 'valid-response.xml')) as fp:
+    VALID_DOC = etree.fromstring(fp.read().encode('utf-8'))
+
+with open(path.join(here, 'too-many-commas.xml')) as fp:
+    TOO_MANY_COMMAS_DOC = etree.fromstring(fp.read().encode('utf-8'))
+
+VALID_ROLE_ARN = "arn:aws:iam::123456789012:role/admin"
+
+def test_parsing_valid_response():
+    roles = aws_google_auth.parse_roles(VALID_DOC)
+    assert(VALID_ROLE_ARN in roles)
+
+def test_parsing_extra_commas():
+    roles = aws_google_auth.parse_roles(TOO_MANY_COMMAS_DOC)
+    assert(VALID_ROLE_ARN in roles)
+
diff --git a/aws_google_auth/tests/too-many-commas.xml b/aws_google_auth/tests/too-many-commas.xml
@@ -0,0 +1,79 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<saml2p:Response xmlns:saml2p="urn:oasis:names:tc:SAML:2.0:protocol" Destination="https://signin.aws.amazon.com/saml" ID="_7c434be06bf79a781dae9e7ed0024679" IssueInstant="2017-07-24T10:31:41.125Z" Version="2.0">
+  <saml2:Issuer xmlns:saml2="urn:oasis:names:tc:SAML:2.0:assertion">https://accounts.google.com/o/saml2?idpid=abcd12345</saml2:Issuer>
+  <saml2p:Status>
+    <saml2p:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
+  </saml2p:Status>
+  <saml2:Assertion xmlns:saml2="urn:oasis:names:tc:SAML:2.0:assertion" ID="_b1dd2c0469d905dfb1e10751d6feae95" IssueInstant="2017-07-24T10:31:41.125Z" Version="2.0">
+    <saml2:Issuer>https://accounts.google.com/o/saml2?idpid=abcd12345</saml2:Issuer>
+    <ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
+      <ds:SignedInfo>
+        <ds:CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+        <ds:SignatureMethod Algorithm="http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"/>
+        <ds:Reference URI="#_b1dd2c0469d905dfb1e10751d6feae95">
+          <ds:Transforms>
+            <ds:Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
+            <ds:Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+          </ds:Transforms>
+          <ds:DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/>
+          <ds:DigestValue>GbaJHVPpMT7JJEn+DtohU/tzd5b/BiZ9+It3sd2LB5Y=</ds:DigestValue>
+        </ds:Reference>
+      </ds:SignedInfo>
+      <ds:SignatureValue>dJxZmFNw+rY07AV7Ex1Kbvn9ZiGE4VKwYELwxkrejgEiVeAteyaw8rQfeHDF1UhZJ/2JTHWs3uk+
+VoWZcI1qcWO3HRjZ/jz7DXH/QGVIBYe447sr9o2RC2WfpjAYTDJ5rN5nPmrQKXxREfFzsZXJutcj
+iPGXDNCC4SsWmKDaqbpWiDKhw+wRxtGxEXB2Ny11dRL6sCIHCdq86H55EXcq2YqL5I/ryMcWt3L0
+SZ5B9aq80omhear/24M1HyL35dmxVUFODrYBxMQ+7Lw6/XUCA2k60MjcsHQW+BJZGwFJBL0HJywu
+bc10BKTA89jbXyBtdoagtWRhF6LJzjL5bImLGA==</ds:SignatureValue>
+      <ds:KeyInfo>
+        <ds:X509Data>
+          <ds:X509SubjectName>ST=California,C=US,OU=Google For Work,CN=Google,L=Mountain View,O=Google Inc.</ds:X509SubjectName>
+          <ds:X509Certificate>MIIDdDCCAlygAwIBAgIGAVXC/OcnMA0GCSqGSIb3DQEBCwUAMHsxFDASBgNVBAoTC0dvb2dsZSBJ
+bmMuMRYwFAYDVQQHEw1Nb3VudGFpbiBWaWV3MQ8wDQYDVQQDEwZHb29nbGUxGDAWBgNVBAsTD0dv
+b2dsZSBGb3IgV29yazELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWEwHhcNMTYwNzA3
+MDEzMzE5WhcNMjEwNzA2MDEzMzE5WjB7MRQwEgYDVQQKEwtHb29nbGUgSW5jLjEWMBQGA1UEBxMN
+TW91bnRhaW4gVmlldzEPMA0GA1UEAxMGR29vZ2xlMRgwFgYDVQQLEw9Hb29nbGUgRm9yIFdvcmsx
+CzAJBgNVBAYTAlVTMRMwEQYDVQQIEwpDYWxpZm9ybmlhMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
+MIIBCgKCAQEAhkv0Sr7ALfc58YrnLXzVGfTRg1T9xUfuZqhdu80BgHTfaJDLX66icHHRRoso/hho
+EIYo1pUQTq0DtgmqkLg9rAup3rR+pImfcHBC55+vMDoEf5t88H/i0qDn3r63PxeULRoFIkCX9aVG
+uUPDe2CHAxB1UXUxyDf7ZAdIQJLPJdOQlsNRleBBoek4vuo2ZHv+A2tbAhE8/rIoQlDvXSpCZ9P7
+m9TrFOb7tB4pHjJjESdmqcnEFc5zepAT8IuRAGZ1OkjJs74JUp+03do8scTMXzvVi4jefpyXhnoN
+C0da4OwPig7UmbDsrSCGbqz29UgxmGUmSnLchpkglw1eET5hTwIDAQABMA0GCSqGSIb3DQEBCwUA
+A4IBAQAA5WBtCPlaSIm1NIpKYd2x8qfeKc2YsxbAPukgUFaRDl1uxGw1HdzNzUp9X4JOF/futpw/
+yhmw9o1GHBukIdj0mJRt8O9szRdkJmx4EfbY5bTVzkQ7QGv9FI1LBD6z6KgJEOxEGpDbh2Z8uyW8
+HvxXgZgiyan53FauVJe+UuAkBy2ynJcVKK3+vUEISFXn1oh5SPOmi+2R4WKSgyTqOKpuowHHHg9u
+EbwwnXPMU4q3QLG1oDrp0ZvVuprvJaoWd5zIt/TYB3Hb5oEO7Imwx1n9K9QskYmFygR9rdJ6VS7L
+6/h6rcL/dKjm4pU0Dgk9h9Hi8ps7Mn+nRRhsWQbiD59n</ds:X509Certificate>
+        </ds:X509Data>
+      </ds:KeyInfo>
+    </ds:Signature>
+    <saml2:Subject>
+      <saml2:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:unspecified">[email protected]</saml2:NameID>
+      <saml2:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
+        <saml2:SubjectConfirmationData NotOnOrAfter="2017-07-24T10:36:41.125Z" Recipient="https://signin.aws.amazon.com/saml"/>
+      </saml2:SubjectConfirmation>
+    </saml2:Subject>
+    <saml2:Conditions NotBefore="2017-07-24T10:26:41.125Z" NotOnOrAfter="2017-07-24T10:36:41.125Z">
+      <saml2:AudienceRestriction>
+        <saml2:Audience>https://signin.aws.amazon.com/saml</saml2:Audience>
+      </saml2:AudienceRestriction>
+    </saml2:Conditions>
+    <saml2:AttributeStatement>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/RoleSessionName">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">[email protected]</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/Role">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/admin,arn:aws:iam::123456789012:saml-provider/GoogleApps,</saml2:AttributeValue>
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/read-only,arn:aws:iam::123456789012:saml-provider/GoogleApps,</saml2:AttributeValue>
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/test,arn:aws:iam::123456789012:saml-provider/GoogleApps,</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/SessionDuration">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">28800</saml2:AttributeValue>
+      </saml2:Attribute>
+    </saml2:AttributeStatement>
+    <saml2:AuthnStatement AuthnInstant="2017-07-24T10:31:38.000Z" SessionIndex="_b1dd2c0469d905dfb1e10751d6feae95">
+      <saml2:AuthnContext>
+        <saml2:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:unspecified</saml2:AuthnContextClassRef>
+      </saml2:AuthnContext>
+    </saml2:AuthnStatement>
+  </saml2:Assertion>
+</saml2p:Response>
diff --git a/aws_google_auth/tests/valid-response.xml b/aws_google_auth/tests/valid-response.xml
@@ -0,0 +1,79 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<saml2p:Response xmlns:saml2p="urn:oasis:names:tc:SAML:2.0:protocol" Destination="https://signin.aws.amazon.com/saml" ID="_7c434be06bf79a781dae9e7ed0024679" IssueInstant="2017-07-24T10:31:41.125Z" Version="2.0">
+  <saml2:Issuer xmlns:saml2="urn:oasis:names:tc:SAML:2.0:assertion">https://accounts.google.com/o/saml2?idpid=abcd12345</saml2:Issuer>
+  <saml2p:Status>
+    <saml2p:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
+  </saml2p:Status>
+  <saml2:Assertion xmlns:saml2="urn:oasis:names:tc:SAML:2.0:assertion" ID="_b1dd2c0469d905dfb1e10751d6feae95" IssueInstant="2017-07-24T10:31:41.125Z" Version="2.0">
+    <saml2:Issuer>https://accounts.google.com/o/saml2?idpid=abcd12345</saml2:Issuer>
+    <ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
+      <ds:SignedInfo>
+        <ds:CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+        <ds:SignatureMethod Algorithm="http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"/>
+        <ds:Reference URI="#_b1dd2c0469d905dfb1e10751d6feae95">
+          <ds:Transforms>
+            <ds:Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
+            <ds:Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
+          </ds:Transforms>
+          <ds:DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/>
+          <ds:DigestValue>GbaJHVPpMT7JJEn+DtohU/tzd5b/BiZ9+It3sd2LB5Y=</ds:DigestValue>
+        </ds:Reference>
+      </ds:SignedInfo>
+      <ds:SignatureValue>dJxZmFNw+rY07AV7Ex1Kbvn9ZiGE4VKwYELwxkrejgEiVeAteyaw8rQfeHDF1UhZJ/2JTHWs3uk+
+VoWZcI1qcWO3HRjZ/jz7DXH/QGVIBYe447sr9o2RC2WfpjAYTDJ5rN5nPmrQKXxREfFzsZXJutcj
+iPGXDNCC4SsWmKDaqbpWiDKhw+wRxtGxEXB2Ny11dRL6sCIHCdq86H55EXcq2YqL5I/ryMcWt3L0
+SZ5B9aq80omhear/24M1HyL35dmxVUFODrYBxMQ+7Lw6/XUCA2k60MjcsHQW+BJZGwFJBL0HJywu
+bc10BKTA89jbXyBtdoagtWRhF6LJzjL5bImLGA==</ds:SignatureValue>
+      <ds:KeyInfo>
+        <ds:X509Data>
+          <ds:X509SubjectName>ST=California,C=US,OU=Google For Work,CN=Google,L=Mountain View,O=Google Inc.</ds:X509SubjectName>
+          <ds:X509Certificate>MIIDdDCCAlygAwIBAgIGAVXC/OcnMA0GCSqGSIb3DQEBCwUAMHsxFDASBgNVBAoTC0dvb2dsZSBJ
+bmMuMRYwFAYDVQQHEw1Nb3VudGFpbiBWaWV3MQ8wDQYDVQQDEwZHb29nbGUxGDAWBgNVBAsTD0dv
+b2dsZSBGb3IgV29yazELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWEwHhcNMTYwNzA3
+MDEzMzE5WhcNMjEwNzA2MDEzMzE5WjB7MRQwEgYDVQQKEwtHb29nbGUgSW5jLjEWMBQGA1UEBxMN
+TW91bnRhaW4gVmlldzEPMA0GA1UEAxMGR29vZ2xlMRgwFgYDVQQLEw9Hb29nbGUgRm9yIFdvcmsx
+CzAJBgNVBAYTAlVTMRMwEQYDVQQIEwpDYWxpZm9ybmlhMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
+MIIBCgKCAQEAhkv0Sr7ALfc58YrnLXzVGfTRg1T9xUfuZqhdu80BgHTfaJDLX66icHHRRoso/hho
+EIYo1pUQTq0DtgmqkLg9rAup3rR+pImfcHBC55+vMDoEf5t88H/i0qDn3r63PxeULRoFIkCX9aVG
+uUPDe2CHAxB1UXUxyDf7ZAdIQJLPJdOQlsNRleBBoek4vuo2ZHv+A2tbAhE8/rIoQlDvXSpCZ9P7
+m9TrFOb7tB4pHjJjESdmqcnEFc5zepAT8IuRAGZ1OkjJs74JUp+03do8scTMXzvVi4jefpyXhnoN
+C0da4OwPig7UmbDsrSCGbqz29UgxmGUmSnLchpkglw1eET5hTwIDAQABMA0GCSqGSIb3DQEBCwUA
+A4IBAQAA5WBtCPlaSIm1NIpKYd2x8qfeKc2YsxbAPukgUFaRDl1uxGw1HdzNzUp9X4JOF/futpw/
+yhmw9o1GHBukIdj0mJRt8O9szRdkJmx4EfbY5bTVzkQ7QGv9FI1LBD6z6KgJEOxEGpDbh2Z8uyW8
+HvxXgZgiyan53FauVJe+UuAkBy2ynJcVKK3+vUEISFXn1oh5SPOmi+2R4WKSgyTqOKpuowHHHg9u
+EbwwnXPMU4q3QLG1oDrp0ZvVuprvJaoWd5zIt/TYB3Hb5oEO7Imwx1n9K9QskYmFygR9rdJ6VS7L
+6/h6rcL/dKjm4pU0Dgk9h9Hi8ps7Mn+nRRhsWQbiD59n</ds:X509Certificate>
+        </ds:X509Data>
+      </ds:KeyInfo>
+    </ds:Signature>
+    <saml2:Subject>
+      <saml2:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:unspecified">[email protected]</saml2:NameID>
+      <saml2:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
+        <saml2:SubjectConfirmationData NotOnOrAfter="2017-07-24T10:36:41.125Z" Recipient="https://signin.aws.amazon.com/saml"/>
+      </saml2:SubjectConfirmation>
+    </saml2:Subject>
+    <saml2:Conditions NotBefore="2017-07-24T10:26:41.125Z" NotOnOrAfter="2017-07-24T10:36:41.125Z">
+      <saml2:AudienceRestriction>
+        <saml2:Audience>https://signin.aws.amazon.com/saml</saml2:Audience>
+      </saml2:AudienceRestriction>
+    </saml2:Conditions>
+    <saml2:AttributeStatement>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/RoleSessionName">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">[email protected]</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/Role">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/admin,arn:aws:iam::123456789012:saml-provider/GoogleApps</saml2:AttributeValue>
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/read-only,arn:aws:iam::123456789012:saml-provider/GoogleApps</saml2:AttributeValue>
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">arn:aws:iam::123456789012:role/test,arn:aws:iam::123456789012:saml-provider/GoogleApps</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/SessionDuration">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:anyType">28800</saml2:AttributeValue>
+      </saml2:Attribute>
+    </saml2:AttributeStatement>
+    <saml2:AuthnStatement AuthnInstant="2017-07-24T10:31:38.000Z" SessionIndex="_b1dd2c0469d905dfb1e10751d6feae95">
+      <saml2:AuthnContext>
+        <saml2:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:unspecified</saml2:AuthnContextClassRef>
+      </saml2:AuthnContext>
+    </saml2:AuthnStatement>
+  </saml2:Assertion>
+</saml2p:Response>