package aws

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"

	"code.crute.us/mcrute/cloud-identity-broker/app/models"
	"code.crute.us/mcrute/golib/secrets"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/sts"
)

// defaultRegion is the region where AWS currently has may authoritative
// service like IAM and ancillary auth services.
var defaultRegion = aws.String("us-east-1")

type Region struct {
	Name    string
	Enabled bool
}

// AWSClient is a client for working with the AWS APIs.
//
// Instances of AWSClient are safe for concurrent access.
type AWSClient interface {
	AssumeRole(string, *string) (*sts.Credentials, error)
	GetFederationURL(string, string) (string, error)
	GetRegionList() ([]*Region, error)
}

type client struct {
	Credentials                credentials.Value
	ARN                        string
	ConsoleSessionDurationSecs int64
}

var _ AWSClient = (*client)(nil)

// NewAWSClientFromAccount returns a new AWSClient based on an account.
//
// An account is actually more accurately called an assumable role. Each
// account contains a vault material set path which is used to fetch all of the
// credentials for that AWS account and is filtered to one assumable role ARN
// which is used as the scope for this AWS client. Thus even if an account has
// multiple roles there must be one instance of the AWS client per account/role
// pair.
func NewAWSClientFromAccount(ctx context.Context, a *models.Account, sc secrets.Client) (AWSClient, error) {
	u, _, err := sc.AWSIAMUser(ctx, a.AdminVaultMaterial)
	if err != nil {
		return nil, err
	}

	return &client{
		Credentials: credentials.Value{
			AccessKeyID:     u.AccessKeyId,
			SecretAccessKey: u.SecretAccessKey,
		},
		ARN:                        a.AssumedRoleARN,
		ConsoleSessionDurationSecs: a.ConsoleSessionDurationSecs(),
	}, nil
}

// AssumeRole uses an IAM user credential with higher privilege to assume a
// role in an AWS account and region. It returns the STS credentials.
//
// Not all credentials work for all regions. For newer regions and opt-in
// regions AWS has been siloing assumed role credentials to that region so it's
// important to use the correct regional endpoint to fetch the credentials.
//
// Note that this is not simply a passthrough to Vault's AWS backend
// because the Vault backend will only call AssumeRole within the region
// of the root account configured in Vault and clients of this method
// need to be able to make the correct calls in the region of their
// choice.
func (a *client) AssumeRole(user string, region *string) (*sts.Credentials, error) {
	if region != nil && *region == "global" {
		region = nil
	}

	s, err := session.NewSessionWithOptions(session.Options{
		Config: aws.Config{
			Region:      region,
			Credentials: credentials.NewStaticCredentialsFromCreds(a.Credentials),
		},
		SharedConfigState: session.SharedConfigDisable,
	})
	if err != nil {
		return nil, err
	}

	result, err := sts.New(s).AssumeRole(&sts.AssumeRoleInput{
		RoleArn:         aws.String(a.ARN),
		RoleSessionName: aws.String(user),
		DurationSeconds: aws.Int64(a.ConsoleSessionDurationSecs),
	})
	if err != nil {
		fmt.Printf("AWS AssumeRole Error: %s\n", err.(awserr.Error).Message())
		return nil, err
	}

	return result.Credentials, nil
}

// GetRegionList returns all of the EC2 regions that are enabled for an
// account.
//
// All regions should return the same answer to this query so just send it to
// the default region.
func (a *client) GetRegionList() ([]*Region, error) {
	r, err := a.AssumeRole("region-list", defaultRegion)
	if err != nil {
		return nil, err
	}

	s, err := session.NewSessionWithOptions(session.Options{
		Config: aws.Config{
			Region: defaultRegion,
			Credentials: credentials.NewStaticCredentials(
				*r.AccessKeyId,
				*r.SecretAccessKey,
				*r.SessionToken,
			),
		},
		SharedConfigState: session.SharedConfigDisable,
	})
	if err != nil {
		return nil, err
	}

	ro, err := ec2.New(s).DescribeRegions(&ec2.DescribeRegionsInput{
		AllRegions: aws.Bool(true),
	})
	if err != nil {
		return nil, err
	}

	out := []*Region{}
	for _, r := range ro.Regions {
		out = append(out, &Region{
			Name:    *r.RegionName,
			Enabled: (*r.OptInStatus != "not-opted-in"),
		})
	}

	return out, nil
}

// GetFederationURL assumes a role and returns a URL that can be used to login
// to the AWS console for that role.
//
// This URL can be used for any region but the endpoints to return the URL only
// work in us-east-1.
func (a *client) GetFederationURL(user string, issuerEndpoint string) (string, error) {
	r, err := a.AssumeRole(user, defaultRegion)
	if err != nil {
		return "", err
	}

	sp, _ := json.Marshal(map[string]string{
		"sessionId":    *r.AccessKeyId,
		"sessionKey":   *r.SecretAccessKey,
		"sessionToken": *r.SessionToken,
	})

	client := &http.Client{}
	fedResp, err := client.Do(&http.Request{
		Method: http.MethodGet,
		URL: &url.URL{
			Scheme: "https",
			Host:   "signin.aws.amazon.com",
			Path:   "/federation",
			RawQuery: url.Values{
				"Action":          []string{"getSigninToken"},
				"Session":         []string{string(sp)},
				"SessionDuration": []string{strconv.FormatInt(a.ConsoleSessionDurationSecs, 10)},
			}.Encode(),
		},
	})
	if err != nil {
		return "", err
	}

	if fedResp.StatusCode != 200 {
		return "", fmt.Errorf("federation endpoint returned HTTP %d", fedResp.StatusCode)
	}

	body, err := ioutil.ReadAll(fedResp.Body)
	if err != nil {
		return "", err
	}

	var fedData map[string]string
	if err = json.Unmarshal(body, &fedData); err != nil {
		return "", err
	}

	u, _ := url.Parse("https://signin.aws.amazon.com/federation")
	u.RawQuery = url.Values{
		"Action":      []string{"login"},
		"Issuer":      []string{issuerEndpoint},
		"Destination": []string{"https://console.aws.amazon.com/"},
		"SigninToken": []string{fedData["SigninToken"]},
	}.Encode()

	return u.String(), nil
}