Support saml2tencentcloud

cloud_account.go

@@ -0,0 +1,136 @@
+package saml2aws
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+
+	"github.com/PuerkitoBio/goquery"
+	"github.com/pkg/errors"
+	"github.com/versent/saml2aws/v2/pkg/cloud"
+)
+
+// CloudAccount holds the AWS account name and roles
+type CloudAccount struct {
+	Name  string
+	Roles []*CloudRole
+}
+
+func AssignAWSAccounts(awsRoles []*CloudRole, samlXml []byte, samlAssertion string) ([]*CloudRole, error) {
+	roleArnMap := make(map[string]*CloudRole)
+	for _, role := range awsRoles {
+		roleArnMap[role.RoleARN] = role
+	}
+
+	aud, err := ExtractDestinationURL(samlXml)
+	if err != nil {
+		return nil, errors.Wrap(err, "Error parsing destination URL.")
+	}
+
+	res, err := http.PostForm(aud, url.Values{"SAMLResponse": {samlAssertion}})
+	if err != nil {
+		return nil, errors.Wrap(err, "Error retrieving cloud SAML login form.")
+	}
+	if res.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("Unexpected status code: %d", res.StatusCode)
+	}
+
+	data, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, errors.Wrap(err, "Error retrieving AWS login body.")
+	}
+
+	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) {
+		name := s.Find("div.saml-account-name").Text()
+		s.Find("label").Each(func(i int, s *goquery.Selection) {
+			arn, _ := s.Attr("for")
+			roleArnMap[arn].Account = name
+			// log.Println("Marked role", arn, "as belonging to account", name)
+		})
+	})
+
+	return awsRoles, nil
+}
+
+// ParseCloudAccounts extract the aws accounts from the saml assertion
+func ParseCloudAccounts(provider cloud.Provider, audience string, samlAssertion string) ([]*CloudAccount, error) {
+
+	switch provider {
+	case cloud.AWS:
+		res, err := http.PostForm(audience, url.Values{"SAMLResponse": {samlAssertion}})
+		if err != nil {
+			return nil, errors.Wrap(err, "error retrieving cloud SAML login form")
+		}
+		if res.StatusCode != http.StatusOK {
+			return nil, fmt.Errorf("unexpected status code: %d", res.StatusCode)
+		}
+
+		data, err := io.ReadAll(res.Body)
+		if err != nil {
+			return nil, errors.Wrap(err, "error retrieving AWS login body")
+		}
+		return ExtractAWSAccounts(data)
+	case cloud.TencentCloud:
+		return nil, nil
+	default:
+		return nil, fmt.Errorf("unsupported cloud provider: %s", provider)
+	}
+}
+
+// ExtractAWSAccounts extract the accounts from the AWS html page
+func ExtractAWSAccounts(data []byte) ([]*CloudAccount, error) {
+	accounts := make([]*CloudAccount, 0)
+
+	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(CloudAccount)
+		account.Name = s.Find("div.saml-account-name").Text()
+		s.Find("label").Each(func(i int, s *goquery.Selection) {
+			role := new(CloudRole)
+			role.Name = s.Text()
+			role.RoleARN, _ = s.Attr("for")
+			account.Roles = append(account.Roles, role)
+		})
+		accounts = append(accounts, account)
+	})
+
+	return accounts, nil
+}
+
+// AssignPrincipals assign principal from roles
+func AssignPrincipals(awsRoles []*CloudRole, cloudAccounts []*CloudAccount) {
+
+	awsPrincipalARNs := make(map[string]string)
+	for _, awsRole := range awsRoles {
+		awsPrincipalARNs[awsRole.RoleARN] = awsRole.PrincipalARN
+	}
+
+	for _, awsAccount := range cloudAccounts {
+		for _, awsRole := range awsAccount.Roles {
+			awsRole.PrincipalARN = awsPrincipalARNs[awsRole.RoleARN]
+		}
+	}
+
+}
+
+// LocateRole locate role by name
+func LocateRole(awsRoles []*CloudRole, roleName string) (*CloudRole, error) {
+	for _, awsRole := range awsRoles {
+		if awsRole.RoleARN == roleName {
+			return awsRole, nil
+		}
+	}
+
+	return nil, fmt.Errorf("Supplied RoleArn not found in saml assertion: %s", roleName)
+}

aws_account_test.go → cloud_account_test.go

@@ -36,16 +36,16 @@ func TestExtractAWSAccounts(t *testing.T) {
 }
 
 func TestAssignPrincipals(t *testing.T) {
-	awsRoles := []*AWSRole{
+	awsRoles := []*CloudRole{
 		{
 			PrincipalARN: "arn:aws:iam::000000000001:saml-provider/test-idp",
 			RoleARN:      "arn:aws:iam::000000000001:role/Development",
 		},
 	}
 
-	awsAccounts := []*AWSAccount{
+	awsAccounts := []*CloudAccount{
 		{
-			Roles: []*AWSRole{
+			Roles: []*CloudRole{
 				{
 					RoleARN: "arn:aws:iam::000000000001:role/Development",
 				},
@@ -59,7 +59,7 @@ func TestAssignPrincipals(t *testing.T) {
 }
 
 func TestLocateRole(t *testing.T) {
-	awsRoles := []*AWSRole{
+	awsRoles := []*CloudRole{
 		{
 			PrincipalARN: "arn:aws:iam::000000000001:saml-provider/test-idp",
 			RoleARN:      "arn:aws:iam::000000000001:role/Development",