package aws

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/hashicorp/errwrap"
	cleanhttp "github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/go-hclog"
	uuid "github.com/hashicorp/go-uuid"
	"github.com/hashicorp/vault/api"
	awsauth "github.com/hashicorp/vault/builtin/credential/aws"
	"github.com/hashicorp/vault/command/agent/auth"
)

const (
	typeEC2          = "ec2"
	typeIAM          = "iam"
	identityEndpoint = "http://169.254.169.254/latest/dynamic/instance-identity"
)

type awsMethod struct {
	logger       hclog.Logger
	authType     string
	nonce        string
	mountPath    string
	role         string
	headerValue  string
	accessKey    string
	secretKey    string
	sessionToken string
}

func NewAWSAuthMethod(conf *auth.AuthConfig) (auth.AuthMethod, error) {
	if conf == nil {
		return nil, errors.New("empty config")
	}
	if conf.Config == nil {
		return nil, errors.New("empty config data")
	}

	a := &awsMethod{
		logger:    conf.Logger,
		mountPath: conf.MountPath,
	}

	typeRaw, ok := conf.Config["type"]
	if !ok {
		return nil, errors.New("missing 'type' value")
	}
	a.authType, ok = typeRaw.(string)
	if !ok {
		return nil, errors.New("could not convert 'type' config value to string")
	}

	roleRaw, ok := conf.Config["role"]
	if !ok {
		return nil, errors.New("missing 'role' value")
	}
	a.role, ok = roleRaw.(string)
	if !ok {
		return nil, errors.New("could not convert 'role' config value to string")
	}

	switch {
	case a.role == "":
		return nil, errors.New("'role' value is empty")
	case a.authType == "":
		return nil, errors.New("'type' value is empty")
	case a.authType != typeEC2 && a.authType != typeIAM:
		return nil, errors.New("'type' value is invalid")
	}

	accessKeyRaw, ok := conf.Config["access_key"]
	if ok {
		a.accessKey, ok = accessKeyRaw.(string)
		if !ok {
			return nil, errors.New("could not convert 'access_key' value into string")
		}
	}

	secretKeyRaw, ok := conf.Config["secret_key"]
	if ok {
		a.secretKey, ok = secretKeyRaw.(string)
		if !ok {
			return nil, errors.New("could not convert 'secret_key' value into string")
		}
	}

	sessionTokenRaw, ok := conf.Config["session_token"]
	if ok {
		a.sessionToken, ok = sessionTokenRaw.(string)
		if !ok {
			return nil, errors.New("could not convert 'session_token' value into string")
		}
	}

	headerValueRaw, ok := conf.Config["header_value"]
	if ok {
		a.headerValue, ok = headerValueRaw.(string)
		if !ok {
			return nil, errors.New("could not convert 'header_value' value into string")
		}
	}

	return a, nil
}

func (a *awsMethod) Authenticate(ctx context.Context, client *api.Client) (retToken string, retData map[string]interface{}, retErr error) {
	a.logger.Trace("beginning authentication")

	data := make(map[string]interface{})

	switch a.authType {
	case typeEC2:
		client := cleanhttp.DefaultClient()

		// Fetch document
		{
			req, err := http.NewRequest("GET", fmt.Sprintf("%s/document", identityEndpoint), nil)
			if err != nil {
				retErr = errwrap.Wrapf("error creating request: {{err}}", err)
				return
			}
			req = req.WithContext(ctx)
			resp, err := client.Do(req)
			if err != nil {
				retErr = errwrap.Wrapf("error fetching instance document: {{err}}", err)
				return
			}
			if resp == nil {
				retErr = errors.New("empty response fetching instance document")
				return
			}
			defer resp.Body.Close()
			doc, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				retErr = errwrap.Wrapf("error reading instance document response body: {{err}}", err)
				return
			}
			data["identity"] = base64.StdEncoding.EncodeToString(doc)
		}

		// Fetch signature
		{
			req, err := http.NewRequest("GET", fmt.Sprintf("%s/signature", identityEndpoint), nil)
			if err != nil {
				retErr = errwrap.Wrapf("error creating request: {{err}}", err)
				return
			}
			req = req.WithContext(ctx)
			resp, err := client.Do(req)
			if err != nil {
				retErr = errwrap.Wrapf("error fetching instance document signature: {{err}}", err)
				return
			}
			if resp == nil {
				retErr = errors.New("empty response fetching instance document signature")
				return
			}
			defer resp.Body.Close()
			sig, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				retErr = errwrap.Wrapf("error reading instance document signature response body: {{err}}", err)
				return
			}
			data["signature"] = string(sig)
		}

		// Add the reauthentication value, if we have one
		if a.nonce == "" {
			uuid, err := uuid.GenerateUUID()
			if err != nil {
				retErr = errwrap.Wrapf("error generating uuid for reauthentication value: {{err}}", err)
				return
			}
			a.nonce = uuid
		}
		data["nonce"] = a.nonce

	default:
		var err error
		data, err = awsauth.GenerateLoginData(a.accessKey, a.secretKey, a.sessionToken, a.headerValue)
		if err != nil {
			retErr = errwrap.Wrapf("error creating login value: {{err}}", err)
			return
		}
	}

	data["role"] = a.role

	return fmt.Sprintf("%s/login", a.mountPath), data, nil
}

func (a *awsMethod) NewCreds() chan struct{} {
	return nil
}

func (a *awsMethod) CredSuccess() {
}

func (a *awsMethod) Shutdown() {
}
