Sync flag

pinko-fowle · Aug 24, 2018 · dac4553 · dac4553
1 parent cdb739b
commit dac4553
Show file tree

Hide file tree

Showing 5 changed files with 206 additions and 67 deletions.
diff --git a/src/aws/role.rs b/src/aws/role.rs
@@ -60,7 +60,7 @@ pub fn assume_role(
     let provider = StaticProvider::new_minimal(String::from(""), String::from(""));
     let client = StsClient::new(RequestDispatcher::default(), provider, Region::default());
 
-    error!("Assuming role: {:?}", &req);
+    trace!("Assuming role: {:?}", &req);
 
     client
         .assume_role_with_saml(&req)

diff --git a/src/main.rs b/src/main.rs
@@ -82,6 +82,10 @@ pub struct Opt {
     /// Silence all output
     #[structopt(short = "q", long = "quiet")]
     pub quiet: bool,
+
+    /// Run in a synchronous manner (no parallel)
+    #[structopt(short = "s", long = "sync")]
+    pub synchronous: bool,
 }
 
 fn main() -> Result<(), Error> {
@@ -97,13 +101,13 @@ fn main() -> Result<(), Error> {
         .format(flexi_logger::with_thread)
         .start()?;*/
 
-    /*let log_level = match opt.verbosity {
+    let log_level = match opt.verbosity {
         0 => "info",
         1 => "debug",
-        _ => "info",
+        _ => "trace",
     };
 
-    env::set_var("RUST_LOG", format!("{}={}", module_path!(), log_level));*/
+    env::set_var("RUST_LOG", format!("{}={}", module_path!(), log_level));
     pretty_env_logger::init();
 
     let config = Config::new()?;
@@ -124,11 +128,13 @@ fn main() -> Result<(), Error> {
             let mut okta_client = organization.okta_client();
             let (username, password) = organization.credentials(opt.force_new)?;
 
-            let session_token =
-                okta_client.get_session_token(&LoginRequest::from_credentials(username, password))?;
+            let session_token = okta_client.get_session_token(&LoginRequest::from_credentials(
+                username.clone(),
+                password.clone(),
+            ))?;
 
             let session_id = okta_client.get_session_id(&session_token)?;
-            okta_client.set_session_id(session_id);
+            okta_client.set_session_id(session_id.clone());
 
             let profiles = organization
                 .profiles(&okta_client)?
@@ -141,56 +147,71 @@ fn main() -> Result<(), Error> {
                     opt.profiles, organization.name
                 );
             } else {
-                profiles
-                    .par_iter()
-                    .try_for_each(|profile: &Profile| -> Result<(), Error> {
-                        info!("Generating tokens for {}", &profile.id);
-
-                        let saml =
-                            okta_client.get_saml_response(profile.application.link_url.clone())?;
-
-                        trace!("SAML assertion: {:?}", saml);
-
-                        let roles = saml.roles;
-
-                        debug!("SAML Roles: {:?}", &roles);
-
-                        match roles
-                            .into_iter()
-                            .find(|r| r.role_name().map(|r| r == profile.role).unwrap_or(false))
-                        {
-                            Some(role) => {
-                                debug!("Found role: {} in {}", role.role_arn, &profile.id);
-
-                                let raw_saml = saml.raw;
-
-                                let assumption_response =
-                                    match aws::role::assume_role(role, raw_saml.clone()) {
-                                        Ok(res) => res,
-                                        Err(e) => {
-                                            error!("Erroring, tried SAML: {:?}", &raw_saml);
-                                            bail!("{}, in profile: {:?}", e, profile.id);
-                                        }
-                                    };
-
-                                if let Some(credentials) = assumption_response.credentials {
-                                    debug!("Credentials: {:?}", credentials);
-
-                                    credentials_store
-                                        .lock()
-                                        .unwrap()
-                                        .set_profile(profile.id.clone(), credentials)
-                                } else {
-                                    bail!("Error fetching credentials from assumed AWS role")
+                let credentials_generator = |profile: &Profile| -> Result<(), Error> {
+                    info!("Generating tokens for {}", &profile.id);
+
+                    let mut okta_client = organization.okta_client();
+                    okta_client.set_session_id(session_id.clone());
+
+                    let saml =
+                        match okta_client.get_saml_response(profile.application.link_url.clone()) {
+                            Ok(saml) => saml,
+                            Err(e) => bail!(
+                                "Error getting SAML response for profile {} ({})",
+                                profile.id,
+                                e
+                            ),
+                        };
+
+                    trace!("SAML assertion: {:?}", saml);
+
+                    let roles = saml.roles;
+
+                    debug!("SAML Roles: {:?}", &roles);
+
+                    match roles
+                        .into_iter()
+                        .find(|r| r.role_name().map(|r| r == profile.role).unwrap_or(false))
+                    {
+                        Some(role) => {
+                            debug!("Found role: {} in {}", role.role_arn, &profile.id);
+
+                            let raw_saml = saml.raw;
+
+                            let assumption_response = match aws::role::assume_role(
+                                role,
+                                raw_saml.clone(),
+                            ) {
+                                Ok(res) => res,
+                                Err(e) => {
+                                    bail!("Error assuming role for profile {} ({})", profile.id, e);
                                 }
+                            };
+
+                            if let Some(credentials) = assumption_response.credentials {
+                                debug!("Credentials: {:?}", credentials);
+
+                                credentials_store
+                                    .lock()
+                                    .unwrap()
+                                    .set_profile(profile.id.clone(), credentials)
+                            } else {
+                                bail!("Error fetching credentials from assumed AWS role")
                             }
-                            None => bail!(
-                                "No matching role ({}) found for profile {}",
-                                profile.role,
-                                &profile.id
-                            ),
                         }
-                    })?;
+                        None => bail!(
+                            "No matching role ({}) found for profile {}",
+                            profile.role,
+                            &profile.id
+                        ),
+                    }
+                };
+
+                if opt.synchronous {
+                    profiles.iter().try_for_each(credentials_generator)?
+                } else {
+                    profiles.par_iter().try_for_each(credentials_generator)?
+                }
             }
         }
 

diff --git a/src/okta/mod.rs b/src/okta/mod.rs
@@ -2,7 +2,7 @@ pub mod auth;
 pub mod client;
 pub mod factor;
 
-use failure::Error;
+use failure::{Compat, Error};
 use kuchiki;
 use kuchiki::traits::TendrilSink;
 use okta::auth::LoginRequest;
@@ -106,15 +106,17 @@ impl OktaClient {
 
         trace!("SAML response doc for app {:?}: {}", &app_url, &response);
 
-        if let Ok(saml) = SamlResponse::from_html(response.clone()) {
-            return Ok(saml);
-        } else {
-            debug!("No SAML found for app {:?}, will re-login", &app_url);
+        match SamlResponse::from_html(response.clone()) {
+            Err(SamlResponseError::NotFound) => {
+                debug!("No SAML found for app {:?}, will re-login", &app_url);
 
-            let state_token = extract_state_token(&response)?;
-            let session_token =
-                self.get_session_token(&LoginRequest::from_state_token(state_token))?;
-            self.get_saml_response(app_url)
+                let state_token = extract_state_token(&response)?;
+                let session_token =
+                    self.get_session_token(&LoginRequest::from_state_token(state_token))?;
+                self.get_saml_response(app_url)
+            }
+            Err(e) => Err(e.into()),
+            Ok(saml) => Ok(saml),
         }
     }
 }
@@ -131,18 +133,32 @@ fn extract_state_token(text: &str) -> Result<String, Error> {
 }
 
 impl SamlResponse {
-    pub fn from_html(text: String) -> Result<Self, Error> {
+    pub fn from_html(text: String) -> Result<Self, SamlResponseError> {
         let doc = kuchiki::parse_html().one(text);
 
         if let Some(input_node) = doc.select("input[name='SAMLResponse']").unwrap().next() {
             if let Some(saml) = input_node.attributes.borrow().get("value") {
                 trace!("SAML: {}", saml);
-                saml.parse()
+                saml.parse().map_err(|e: Error| e.into())
             } else {
-                bail!("No `value` found in SAML block")
+                Err(SamlResponseError::NotFound)
             }
         } else {
-            bail!("No SAML block found")
+            Err(SamlResponseError::NotFound)
         }
     }
 }
+
+#[derive(Fail, Debug)]
+pub enum SamlResponseError {
+    #[fail(display = "No SAML found")]
+    NotFound,
+    #[fail(display = "{}", _0)]
+    Invalid(#[cause] Compat<Error>),
+}
+
+impl From<Error> for SamlResponseError {
+    fn from(e: Error) -> SamlResponseError {
+        SamlResponseError::Invalid(e.compat())
+    }
+}
diff --git a/src/saml/mod.rs b/src/saml/mod.rs
@@ -80,4 +80,23 @@ mod tests {
 
         assert_eq!(response.roles, expected_roles);
     }
+
+    #[test]
+    fn parse_response_invalid_no_role() {
+        let mut f =
+            File::open("tests/fixtures/saml_response_invalid_no_role.xml").expect("file not found");
+
+        let mut saml_xml = String::new();
+        f.read_to_string(&mut saml_xml)
+            .expect("something went wrong reading the file");
+
+        let saml_base64 = encode(&saml_xml);
+
+        let response: Error = saml_base64.parse::<Response>().unwrap_err();
+
+        assert_eq!(
+            response.to_string(),
+            "Not enough elements in arn:aws:iam::123456789012:saml-provider/okta-idp"
+        );
+    }
 }
diff --git a/tests/fixtures/saml_response_invalid_no_role.xml b/tests/fixtures/saml_response_invalid_no_role.xml
@@ -0,0 +1,83 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<!-- The XML below is taken from the
+     'SAML Response with Signed Message & Assertion' example SAML response at
+     https://www.samltool.com/generic_sso_res.php,
+     then amended with AWS-specific attributes -->
+
+<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" ID="pfx9970aa50-d74b-e2e1-d463-1dff95e0c9c1" Version="2.0" IssueInstant="2014-07-17T01:01:48Z" Destination="http://sp.example.com/demo1/index.php?acs" InResponseTo="ONELOGIN_4fee3b046395c4e751011e97f8900b5273d56685">
+  <saml:Issuer>http://idp.example.com/metadata.php</saml: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/2000/09/xmldsig#rsa-sha1"/>
+      <ds:Reference URI="#pfx9970aa50-d74b-e2e1-d463-1dff95e0c9c1">
+        <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/2000/09/xmldsig#sha1"/>
+        <ds:DigestValue>tf7oPpgBqPFzCkwWkv6m/vsmJaU=</ds:DigestValue>
+      </ds:Reference>
+    </ds:SignedInfo>
+    <ds:SignatureValue>UYO1rUre+g5Iu+zKjVf6KyP+T200P6GJVqZthXTnQSLr2C3TOQO5X/OAlFFH53jyAfUXajDLMMdeMPlDgL632p9Ejndx7pqe22fzxO4EVRiWXQ5X0HegJ/1OM20w+HHyNhpN43J8hbMxS2xXOOgEcZHcp5A21FPgIAisMYLsXFg=</ds:SignatureValue>
+    <ds:KeyInfo>
+      <ds:X509Data>
+        <ds:X509Certificate>MIICajCCAdOgAwIBAgIBADANBgkqhkiG9w0BAQ0FADBSMQswCQYDVQQGEwJ1czETMBEGA1UECAwKQ2FsaWZvcm5pYTEVMBMGA1UECgwMT25lbG9naW4gSW5jMRcwFQYDVQQDDA5zcC5leGFtcGxlLmNvbTAeFw0xNDA3MTcxNDEyNTZaFw0xNTA3MTcxNDEyNTZaMFIxCzAJBgNVBAYTAnVzMRMwEQYDVQQIDApDYWxpZm9ybmlhMRUwEwYDVQQKDAxPbmVsb2dpbiBJbmMxFzAVBgNVBAMMDnNwLmV4YW1wbGUuY29tMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDZx+ON4IUoIWxgukTb1tOiX3bMYzYQiwWPUNMp+Fq82xoNogso2bykZG0yiJm5o8zv/sd6pGouayMgkx/2FSOdc36T0jGbCHuRSbtia0PEzNIRtmViMrt3AeoWBidRXmZsxCNLwgIV6dn2WpuE5Az0bHgpZnQxTKFek0BMKU/d8wIDAQABo1AwTjAdBgNVHQ4EFgQUGHxYqZYyX7cTxKVODVgZwSTdCnwwHwYDVR0jBBgwFoAUGHxYqZYyX7cTxKVODVgZwSTdCnwwDAYDVR0TBAUwAwEB/zANBgkqhkiG9w0BAQ0FAAOBgQByFOl+hMFICbd3DJfnp2Rgd/dqttsZG/tyhILWvErbio/DEe98mXpowhTkC04ENprOyXi7ZbUqiicF89uAGyt1oqgTUCD1VsLahqIcmrzgumNyTwLGWo17WDAa1/usDhetWAMhgzF/Cnf5ek0nK00m0YZGyc4LzgD0CROMASTWNg==</ds:X509Certificate>
+      </ds:X509Data>
+    </ds:KeyInfo>
+  </ds:Signature>
+  <samlp:Status>
+    <samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
+  </samlp:Status>
+  <saml:Assertion xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:xs="http://www.w3.org/2001/XMLSchema" ID="pfx8dcecb85-a2e5-839b-10f5-7da2be872394" Version="2.0" IssueInstant="2014-07-17T01:01:48Z">
+    <saml:Issuer>http://idp.example.com/metadata.php</saml: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/2000/09/xmldsig#rsa-sha1"/>
+        <ds:Reference URI="#pfx8dcecb85-a2e5-839b-10f5-7da2be872394">
+          <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/2000/09/xmldsig#sha1"/>
+          <ds:DigestValue>2QBPUDYERf3P3loU5ruLKgyJw1Y=</ds:DigestValue>
+        </ds:Reference>
+      </ds:SignedInfo>
+      <ds:SignatureValue>D2eoAgfaHe3HEcgpL8DjbA5MPLrEF+wAotHJG8ku1ej2lPnD96ZUj9b5XIMIAHUgj60Napnrgg3QDfaHgA+ESiOtEx9+yfSULVZZjQLmHaKY8zXoM1KsnWPjsI2yqlYpm1dLu6JiQSnXq7mv6UnHwzTuV67IqCi4/NoX1Kzct84=</ds:SignatureValue>
+      <ds:KeyInfo>
+        <ds:X509Data>
+          <ds:X509Certificate>MIICajCCAdOgAwIBAgIBADANBgkqhkiG9w0BAQ0FADBSMQswCQYDVQQGEwJ1czETMBEGA1UECAwKQ2FsaWZvcm5pYTEVMBMGA1UECgwMT25lbG9naW4gSW5jMRcwFQYDVQQDDA5zcC5leGFtcGxlLmNvbTAeFw0xNDA3MTcxNDEyNTZaFw0xNTA3MTcxNDEyNTZaMFIxCzAJBgNVBAYTAnVzMRMwEQYDVQQIDApDYWxpZm9ybmlhMRUwEwYDVQQKDAxPbmVsb2dpbiBJbmMxFzAVBgNVBAMMDnNwLmV4YW1wbGUuY29tMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDZx+ON4IUoIWxgukTb1tOiX3bMYzYQiwWPUNMp+Fq82xoNogso2bykZG0yiJm5o8zv/sd6pGouayMgkx/2FSOdc36T0jGbCHuRSbtia0PEzNIRtmViMrt3AeoWBidRXmZsxCNLwgIV6dn2WpuE5Az0bHgpZnQxTKFek0BMKU/d8wIDAQABo1AwTjAdBgNVHQ4EFgQUGHxYqZYyX7cTxKVODVgZwSTdCnwwHwYDVR0jBBgwFoAUGHxYqZYyX7cTxKVODVgZwSTdCnwwDAYDVR0TBAUwAwEB/zANBgkqhkiG9w0BAQ0FAAOBgQByFOl+hMFICbd3DJfnp2Rgd/dqttsZG/tyhILWvErbio/DEe98mXpowhTkC04ENprOyXi7ZbUqiicF89uAGyt1oqgTUCD1VsLahqIcmrzgumNyTwLGWo17WDAa1/usDhetWAMhgzF/Cnf5ek0nK00m0YZGyc4LzgD0CROMASTWNg==</ds:X509Certificate>
+        </ds:X509Data>
+      </ds:KeyInfo>
+    </ds:Signature>
+    <saml:Subject>
+      <saml:NameID SPNameQualifier="http://sp.example.com/demo1/metadata.php" Format="urn:oasis:names:tc:SAML:2.0:nameid-format:transient">_ce3d2948b4cf20146dee0a0b3dd6f69b6cf86f62d7</saml:NameID>
+      <saml:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
+        <saml:SubjectConfirmationData NotOnOrAfter="2024-01-18T06:21:48Z" Recipient="http://sp.example.com/demo1/index.php?acs" InResponseTo="ONELOGIN_4fee3b046395c4e751011e97f8900b5273d56685"/>
+      </saml:SubjectConfirmation>
+    </saml:Subject>
+    <saml:Conditions NotBefore="2014-07-17T01:01:18Z" NotOnOrAfter="2024-01-18T06:21:48Z">
+      <saml:AudienceRestriction>
+        <saml:Audience>http://sp.example.com/demo1/metadata.php</saml:Audience>
+      </saml:AudienceRestriction>
+    </saml:Conditions>
+    <saml:AuthnStatement AuthnInstant="2014-07-17T01:01:48Z" SessionNotOnOrAfter="2024-07-17T09:01:48Z" SessionIndex="_be9967abd904ddcae3c0eb4189adbe3f71e327cf93">
+      <saml:AuthnContext>
+        <saml:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:Password</saml:AuthnContextClassRef>
+      </saml:AuthnContext>
+    </saml:AuthnStatement>
+    <saml2:AttributeStatement xmlns:saml2="urn:oasis:names:tc:SAML:2.0:assertion">
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/Role" NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:uri">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:string">arn:aws:iam::123456789012:saml-provider/okta-idp</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/RoleSessionName" NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:basic">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:string">[email protected]</saml2:AttributeValue>
+      </saml2:Attribute>
+      <saml2:Attribute Name="https://aws.amazon.com/SAML/Attributes/SessionDuration" NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:basic">
+        <saml2:AttributeValue xmlns:xs="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:type="xs:string">43200</saml2:AttributeValue>
+      </saml2:Attribute>
+    </saml2:AttributeStatement>
+  </saml:Assertion>
+</samlp:Response>