Repository: Versent/saml2aws
Pull request #22 (2017-01-30 04:09:05)

Account aliases
This addresses #7 by POSTing to the interactive login form provided by AWS.  Account alias and role name is extracted from the markup.  The CLI prompt has been modified to group by account just like the login form.

I also setup `$ make install` to simplify installation.


Patch:
diff --git a/Makefile b/Makefile
index 114289b..7746033 100644
--- a/Makefile
+++ b/Makefile
@@ -21,6 +21,9 @@ compile: deps
 	-output "build/{{.Dir}}_$(VERSION)_{{.OS}}_{{.Arch}}/$(NAME)" \
 	$(shell glide novendor)
 
+install:
+	go install ./cmd/saml2aws
+
 dist:
 	$(eval FILES := $(shell ls build))
 	@rm -rf dist && mkdir dist
diff --git a/README.md b/README.md
index 174f3dc..b7c215e 100644
--- a/README.md
+++ b/README.md
@@ -157,6 +157,12 @@ Then to build the software just run.
 make
 ```
 
+Install the binary to `$GOPATH/bin`.
+
+```
+make install
+```
+
 To release run.
 
 ```
diff --git a/aws_account.go b/aws_account.go
new file mode 100644
index 0000000..ccf3306
--- /dev/null
+++ b/aws_account.go
@@ -0,0 +1,55 @@
+package saml2aws
+
+import (
+	"bytes"
+	"io/ioutil"
+	"net/http"
+	"net/url"
+
+	"github.com/PuerkitoBio/goquery"
+	"github.com/pkg/errors"
+)
+
+type AWSAccount struct {
+	Name  string
+	Roles []*AWSRole
+}
+
+func ParseAWSAccounts(samlAssertion string) ([]*AWSAccount, error) {
+	awsURL := "https://signin.aws.amazon.com/saml"
+
+	res, err := http.PostForm(awsURL, url.Values{"SAMLResponse": {samlAssertion}})
+	if err != nil {
+		return nil, errors.Wrap(err, "error retieving AWS login form")
+	}
+
+	data, err := ioutil.ReadAll(res.Body)
+	if err != nil {
+		return nil, errors.Wrap(err, "error retieving AWS login body")
+	}
+
+	return ExtractAWSAccounts(data)
+}
+
+func ExtractAWSAccounts(data []byte) ([]*AWSAccount, error) {
+	accounts := []*AWSAccount{}
+
+	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(data))
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to build document from response")
+	}
+
+	doc.Find("fieldset > div.saml-account").Each(func(i int, s *goquery.Selection) {
+		account := new(AWSAccount)
+		account.Name = s.Find("div.saml-account-name").Text()
+		s.Find("label").Each(func(i int, s *goquery.Selection) {
+			role := new(AWSRole)
+			role.Name = s.Text()
+			role.RoleARN, _ = s.Attr("for")
+			account.Roles = append(account.Roles, role)
+		})
+		accounts = append(accounts, account)
+	})
+
+	return accounts, nil
+}
diff --git a/aws_role.go b/aws_role.go
index fe9f282..ee4ea52 100644
--- a/aws_role.go
+++ b/aws_role.go
@@ -9,6 +9,7 @@ import (
 type AWSRole struct {
 	RoleARN      string
 	PrincipalARN string
+	Name         string
 }
 
 // ParseAWSRoles parses and splits the roles while also validating the contents
diff --git a/cmd/saml2aws/commands/login.go b/cmd/saml2aws/commands/login.go
index 2c2c9fa..7c1aa3d 100644
--- a/cmd/saml2aws/commands/login.go
+++ b/cmd/saml2aws/commands/login.go
@@ -76,9 +76,33 @@ func Login(profile, providerName string, skipVerify bool) error {
 		return errors.Wrap(err, "error parsing aws roles")
 	}
 
-	role, err := saml2aws.PromptForAWSRoleSelection(awsRoles)
-	if err != nil {
-		return errors.Wrap(err, "error selecting role")
+	var role = new(saml2aws.AWSRole)
+
+	if len(awsRoles) == 1 {
+		role = awsRoles[0]
+	} else if len(awsRoles) == 0 {
+		return errors.Wrap(err, "no roles available")
+	} else {
+		awsPrincipalARNs := make(map[string]string)
+		for _, awsRole := range awsRoles {
+			awsPrincipalARNs[awsRole.RoleARN] = awsRole.PrincipalARN
+		}
+
+		awsAccounts, err := saml2aws.ParseAWSAccounts(samlAssertion)
+		if err != nil {
+			return errors.Wrap(err, "error parsing aws role accounts")
+		}
+
+		for _, awsAccount := range awsAccounts {
+			for _, awsRole := range awsAccount.Roles {
+				awsRole.PrincipalARN = awsPrincipalARNs[awsRole.RoleARN]
+			}
+		}
+
+		role, err = saml2aws.PromptForAWSRoleSelection(awsAccounts)
+		if err != nil {
+			return errors.Wrap(err, "error selecting role")
+		}
 	}
 
 	fmt.Println("Selected role:", role.RoleARN)
diff --git a/input.go b/input.go
index 37732b2..8b4869e 100644
--- a/input.go
+++ b/input.go
@@ -34,18 +34,21 @@ func PromptForLoginDetails(username, hostname string) (*LoginDetails, error) {
 }
 
 // PromptForAWSRoleSelection present a list of roles to the user for selection
-func PromptForAWSRoleSelection(roles []*AWSRole) (*AWSRole, error) {
-
-	if len(roles) == 1 {
-		return roles[0], nil
-	}
+func PromptForAWSRoleSelection(accounts []*AWSAccount) (*AWSRole, error) {
 
 	reader := bufio.NewReader(os.Stdin)
 
 	fmt.Println("Please choose the role you would like to assume: ")
 
-	for i, role := range roles {
-		fmt.Println("[", i, "]: ", role.RoleARN)
+	roles := []*AWSRole{}
+
+	for _, account := range accounts {
+		fmt.Println(account.Name)
+		for _, role := range account.Roles {
+			fmt.Println("[", len(roles), "]: ", role.Name)
+			fmt.Println()
+			roles = append(roles, role)
+		}
 	}
 
 	fmt.Print("Selection: ")


Test patch:
diff --git a/aws_account_test.go b/aws_account_test.go
new file mode 100644
index 0000000..d4159f0
--- /dev/null
+++ b/aws_account_test.go
@@ -0,0 +1,36 @@
+package saml2aws
+
+import (
+	"io/ioutil"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExtractAWSAccounts(t *testing.T) {
+	data, err := ioutil.ReadFile("testdata/saml.html")
+	assert.Nil(t, err)
+
+	accounts, err := ExtractAWSAccounts(data)
+	assert.Nil(t, err)
+	assert.Len(t, accounts, 2)
+
+	account := accounts[0]
+	assert.Equal(t, account.Name, "Account: account-alias (000000000001)")
+
+	assert.Len(t, account.Roles, 2)
+	role := account.Roles[0]
+	assert.Equal(t, role.RoleARN, "arn:aws:iam::000000000001:role/Development")
+	assert.Equal(t, role.Name, "Development")
+	role = account.Roles[1]
+	assert.Equal(t, role.RoleARN, "arn:aws:iam::000000000001:role/Production")
+	assert.Equal(t, role.Name, "Production")
+
+	account = accounts[1]
+	assert.Equal(t, account.Name, "Account: 000000000002")
+
+	assert.Len(t, account.Roles, 1)
+	role = account.Roles[0]
+	assert.Equal(t, role.RoleARN, "arn:aws:iam::000000000002:role/Production")
+	assert.Equal(t, role.Name, "Production")
+}
diff --git a/testdata/saml.html b/testdata/saml.html
new file mode 100644
index 0000000..6d5f5a6
--- /dev/null
+++ b/testdata/saml.html
@@ -0,0 +1,115 @@
+<!DOCTYPE html>
+<html>
+<body>
+
+<div id="container">
+  
+  <h1 class="background">Amazon Web Services Login</h1>
+  
+  <div id="content">
+ 
+  <div id="main_error"></div>
+  
+  <form id="saml_form" name="saml_form" action="/saml" method="post">
+          <input type="hidden" name="RelayState" value="" />
+          <input type="hidden" name="SAMLResponse" value="REDACTED" />
+          <input type="hidden" name="name" value="" />
+          <p style="font-size: 16px; padding-left: 20px;">Select a role:</p>
+          <fieldset>
+            <div  class="saml-account"> <div onClick="expandCollapse(0); expandCollapse(1);">
+              <img id="image0" src="/static/image/down.png" valign="middle"></img>
+              <div class="saml-account-name">Account: account-alias (000000000001)</div>
+              </div>
+              <hr style="border: 1px solid #ddd;">
+              <div id="0" class="saml-account" >  
+                <div class="saml-role" onClick="checkRadio(this);">
+                    <input type="radio" name="roleIndex" value="arn:aws:iam::000000000001:role/Development" class="saml-radio" id="arn:aws:iam::000000000001:role/Development" />
+                    <label for="arn:aws:iam::000000000001:role/Development" class="saml-role-description">Development</label>
+                    <span style="clear: both;"></span>
+                </div>
+                 </div>
+              <div id="1" class="saml-account" >  
+                <div class="saml-role" onClick="checkRadio(this);">
+                    <input type="radio" name="roleIndex" value="arn:aws:iam::000000000001:role/Production" class="saml-radio" id="arn:aws:iam::000000000001:role/Production" />
+                    <label for="arn:aws:iam::000000000001:role/Production" class="saml-role-description">Production</label>
+                    <span style="clear: both;"></span>
+                </div>
+                 </div>
+							</div>
+              <div  class="saml-account"> <div onClick="expandCollapse(2);">
+              <img id="image4" src="/static/image/down.png" valign="middle"></img>
+              <div class="saml-account-name">Account: 000000000002</div>
+              </div>
+              <hr style="border: 1px solid #ddd;">
+              <div id="2" class="saml-account" >  
+                <div class="saml-role" onClick="checkRadio(this);">
+                    <input type="radio" name="roleIndex" value="arn:aws:iam::000000000002:role/Production" class="saml-radio" id="arn:aws:iam::000000000002:role/Production" />
+                    <label for="arn:aws:iam::000000000002:role/Production" class="saml-role-description">Production</label>
+                    <span style="clear: both;"></span>
+                </div>
+                 </div></div>
+          </fieldset>  
+          
+          <br>
+          <div class="buttoninput" id="input_signin_button">
+              <a id="signin_button" class="css3button" href="#" alt="Continue" value="Continue">Sign In</a>
+          </div>
+
+  </form>
+  </div>
+
+  <br/>
+  <div id="smallprint" style="text-align: center;">
+  <div class="textinput" style = "text-align: center; font-size: x-small;">
+    
+    <div style="border-top: 1px solid #ccc; margin-bottom: 12px; padding-top: 1em; text-align: center;">
+        <select id="languageSelector">
+        
+            <option value="en_US"  selected >English</option>
+        
+            <option value="de_DE" >Deutsch</option>
+        
+            <option value="es_ES" >Español</option>
+        
+            <option value="fr_FR" >Français</option>
+        
+            <option value="ja" >日本語</option>
+        
+            <option value="pt_BR" >Português</option>
+        
+            <option value="ko_KR" >한국어</option>
+        
+            <option value="zh_CN" >中文(简体)</option>
+        
+        </select>
+        
+        <script type="text/javascript">
+            var currentYear = new Date().getFullYear()
+
+            $(document).ready(function () {
+                $('#languageSelector').bind("change", function () {
+                        document.cookie= ("noflush_locale=" + $('#languageSelector').val()+ "; domain=.amazon.com");
+                        window.location.reload();
+                });
+            });
+        </script>
+    </div>
+    
+    <a href="https://aws.amazon.com/terms/">Terms of Use</a>
+    <a href="https://aws.amazon.com/privacy/">Privacy Policy</a><br/>
+    
+
+    
+        <script type="text/javascript">
+          document.write('&copy; 1996-' + currentYear + ', Amazon Web Services, Inc. or its affiliates.');
+        </script>
+    
+
+  </div>
+</div>
+
+</div>
+  
+</body>
+</html>
+