// © Broadcom. All Rights Reserved.
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
// SPDX-License-Identifier: Apache-2.0

package internal

// The sts/internal package provides the types for invoking the sts.Issue method.
// The sts.Issue and SessionManager LoginByToken methods require an XML signature.
// Unlike the JRE and .NET runtimes, the Go stdlib does not support XML signing.
// We should considering contributing to the goxmldsig package and gosaml2 to meet
// the needs of sts.Issue rather than maintaining this package long term.
// The tricky part of xmldig is the XML canonicalization (C14N), which is responsible
// for most of the make-your-eyes bleed XML formatting in this package.
// C14N is also why some structures use xml.Name without a field tag and methods modify the xml.Name directly,
// though also working around Go's handling of XML namespace prefixes.
// Most of the types in this package were originally generated from the wsdl and hacked up gen/ scripts,
// but have since been modified by hand.

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"log"
	"path"
	"reflect"
	"strings"

	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
	"github.com/vmware/govmomi/vim25/xml"
)

const (
	XSI    = "http://www.w3.org/2001/XMLSchema-instance"
	WSU    = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd"
	DSIG   = "http://www.w3.org/2000/09/xmldsig#"
	SHA256 = "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"
	Time   = "2006-01-02T15:04:05.000Z"
)

// Security is used as soap.Envelope.Header.Security when signing requests.
type Security struct {
	XMLName             xml.Name `xml:"wsse:Security"`
	WSSE                string   `xml:"xmlns:wsse,attr"`
	WSU                 string   `xml:"xmlns:wsu,attr"`
	Timestamp           Timestamp
	BinarySecurityToken *BinarySecurityToken `xml:",omitempty"`
	UsernameToken       *UsernameToken       `xml:",omitempty"`
	Assertion           string               `xml:",innerxml"`
	Signature           *Signature           `xml:",omitempty"`
}

type Timestamp struct {
	XMLName xml.Name `xml:"wsu:Timestamp"`
	NS      string   `xml:"xmlns:wsu,attr"`
	ID      string   `xml:"wsu:Id,attr"`
	Created string   `xml:"wsu:Created"`
	Expires string   `xml:"wsu:Expires"`
}

func (t *Timestamp) C14N() string {
	return Marshal(t)
}

type BinarySecurityToken struct {
	XMLName      xml.Name `xml:"wsse:BinarySecurityToken"`
	EncodingType string   `xml:"EncodingType,attr"`
	ValueType    string   `xml:"ValueType,attr"`
	ID           string   `xml:"wsu:Id,attr"`
	Value        string   `xml:",chardata"`
}

type UsernameToken struct {
	XMLName  xml.Name `xml:"wsse:UsernameToken"`
	Username string   `xml:"wsse:Username"`
	Password string   `xml:"wsse:Password"`
}

type Signature struct {
	XMLName        xml.Name
	NS             string `xml:"xmlns:ds,attr"`
	ID             string `xml:"Id,attr"`
	SignedInfo     SignedInfo
	SignatureValue Value
	KeyInfo        KeyInfo
}

func (s *Signature) C14N() string {
	return fmt.Sprintf(`<ds:Signature xmlns:ds="%s">%s%s%s</ds:Signature>`,
		DSIG, s.SignedInfo.C14N(), s.SignatureValue.C14N(), s.KeyInfo.C14N())
}

type SignedInfo struct {
	XMLName                xml.Name
	NS                     string `xml:"xmlns:ds,attr,omitempty"`
	CanonicalizationMethod Method
	SignatureMethod        Method
	Reference              []Reference
}

func (s SignedInfo) C14N() string {
	ns := "" // empty in ActAs c14n form for example
	if s.NS != "" {
		ns = fmt.Sprintf(` xmlns:ds="%s"`, s.NS)
	}

	c14n := []string{fmt.Sprintf("<ds:SignedInfo%s>", ns)}
	c14n = append(c14n, s.CanonicalizationMethod.C14N(), s.SignatureMethod.C14N())
	for i := range s.Reference {
		c14n = append(c14n, s.Reference[i].C14N())
	}
	c14n = append(c14n, "</ds:SignedInfo>")

	return strings.Join(c14n, "")
}

type Method struct {
	XMLName   xml.Name
	Algorithm string `xml:",attr"`
}

func (m *Method) C14N() string {
	return mkns("ds", m, &m.XMLName)
}

type Value struct {
	XMLName xml.Name
	Value   string `xml:",innerxml"`
}

func (v *Value) C14N() string {
	return mkns("ds", v, &v.XMLName)
}

type Reference struct {
	XMLName      xml.Name
	URI          string `xml:",attr"`
	Transforms   Transforms
	DigestMethod Method
	DigestValue  Value
}

func (r Reference) C14N() string {
	for i := range r.Transforms.Transform {
		t := &r.Transforms.Transform[i]
		t.XMLName.Local = "ds:Transform"
		t.XMLName.Space = ""

		if t.InclusiveNamespaces != nil {
			name := &t.InclusiveNamespaces.XMLName
			if !strings.HasPrefix(name.Local, "ec:") {
				name.Local = "ec:" + name.Local
				name.Space = ""
			}
			t.InclusiveNamespaces.NS = t.Algorithm
		}
	}

	c14n := []string{
		fmt.Sprintf(`<ds:Reference URI="%s">`, r.URI),
		r.Transforms.C14N(),
		r.DigestMethod.C14N(),
		r.DigestValue.C14N(),
		"</ds:Reference>",
	}

	return strings.Join(c14n, "")
}

func NewReference(id string, val string) Reference {
	sum := sha256.Sum256([]byte(val))

	return Reference{
		XMLName: xml.Name{Local: "ds:Reference"},
		URI:     "#" + id,
		Transforms: Transforms{
			XMLName: xml.Name{Local: "ds:Transforms"},
			Transform: []Transform{
				Transform{
					XMLName:   xml.Name{Local: "ds:Transform"},
					Algorithm: "http://www.w3.org/2001/10/xml-exc-c14n#",
				},
			},
		},
		DigestMethod: Method{
			XMLName:   xml.Name{Local: "ds:DigestMethod"},
			Algorithm: "http://www.w3.org/2001/04/xmlenc#sha256",
		},
		DigestValue: Value{
			XMLName: xml.Name{Local: "ds:DigestValue"},
			Value:   base64.StdEncoding.EncodeToString(sum[:]),
		},
	}
}

type Transforms struct {
	XMLName   xml.Name
	Transform []Transform
}

func (t *Transforms) C14N() string {
	return mkns("ds", t, &t.XMLName)
}

type Transform struct {
	XMLName             xml.Name
	Algorithm           string               `xml:",attr"`
	InclusiveNamespaces *InclusiveNamespaces `xml:",omitempty"`
}

type InclusiveNamespaces struct {
	XMLName    xml.Name
	NS         string `xml:"xmlns:ec,attr,omitempty"`
	PrefixList string `xml:",attr"`
}

type X509Data struct {
	XMLName         xml.Name
	X509Certificate string `xml:",innerxml"`
}

type KeyInfo struct {
	XMLName                xml.Name
	NS                     string                  `xml:"xmlns:ds,attr,omitempty"`
	SecurityTokenReference *SecurityTokenReference `xml:",omitempty"`
	X509Data               *X509Data               `xml:",omitempty"`
}

func (o *KeyInfo) C14N() string {
	names := []*xml.Name{
		&o.XMLName,
	}

	if o.SecurityTokenReference != nil {
		names = append(names, &o.SecurityTokenReference.XMLName)
	}
	if o.X509Data != nil {
		names = append(names, &o.X509Data.XMLName)
	}

	return mkns("ds", o, names...)
}

type SecurityTokenReference struct {
	XMLName       xml.Name           `xml:"wsse:SecurityTokenReference"`
	WSSE11        string             `xml:"xmlns:wsse11,attr,omitempty"`
	TokenType     string             `xml:"wsse11:TokenType,attr,omitempty"`
	Reference     *SecurityReference `xml:",omitempty"`
	KeyIdentifier *KeyIdentifier     `xml:",omitempty"`
}

type SecurityReference struct {
	XMLName   xml.Name `xml:"wsse:Reference"`
	URI       string   `xml:",attr"`
	ValueType string   `xml:",attr"`
}

type KeyIdentifier struct {
	XMLName   xml.Name `xml:"wsse:KeyIdentifier"`
	ID        string   `xml:",innerxml"`
	ValueType string   `xml:",attr"`
}

type Issuer struct {
	XMLName xml.Name
	Format  string `xml:",attr"`
	Value   string `xml:",innerxml"`
}

func (i *Issuer) C14N() string {
	return mkns("saml2", i, &i.XMLName)
}

type Assertion struct {
	XMLName            xml.Name
	ID                 string `xml:",attr"`
	IssueInstant       string `xml:",attr"`
	Version            string `xml:",attr"`
	Issuer             Issuer
	Signature          Signature
	Subject            Subject
	Conditions         Conditions
	AuthnStatement     AuthnStatement
	AttributeStatement AttributeStatement
}

func (a *Assertion) C14N() string {
	start := `<saml2:Assertion xmlns:saml2="%s" ID="%s" IssueInstant="%s" Version="%s">`
	c14n := []string{
		fmt.Sprintf(start, a.XMLName.Space, a.ID, a.IssueInstant, a.Version),
		a.Issuer.C14N(),
		a.Signature.C14N(),
		a.Subject.C14N(),
		a.Conditions.C14N(),
		a.AuthnStatement.C14N(),
		a.AttributeStatement.C14N(),
		`</saml2:Assertion>`,
	}

	return strings.Join(c14n, "")
}

type NameID struct {
	XMLName xml.Name
	Format  string `xml:",attr"`
	ID      string `xml:",innerxml"`
}

type Subject struct {
	XMLName             xml.Name
	NameID              NameID
	SubjectConfirmation SubjectConfirmation
}

func (s *Subject) C14N() string {
	data := &s.SubjectConfirmation.SubjectConfirmationData
	names := []*xml.Name{
		&s.XMLName,
		&s.NameID.XMLName,
		&s.SubjectConfirmation.XMLName,
		&data.XMLName,
	}
	if s.SubjectConfirmation.NameID != nil {
		names = append(names, &s.SubjectConfirmation.NameID.XMLName)
	}
	if data.KeyInfo != nil {
		data.NS = XSI
		data.Type = "saml2:KeyInfoConfirmationDataType"
		data.KeyInfo.XMLName = xml.Name{Local: "ds:KeyInfo"}
		data.KeyInfo.X509Data.XMLName = xml.Name{Local: "ds:X509Data"}
		data.KeyInfo.NS = DSIG
	}
	return mkns("saml2", s, names...)
}

type SubjectConfirmationData struct {
	XMLName      xml.Name
	NS           string `xml:"xmlns:xsi,attr,omitempty"`
	Type         string `xml:"xsi:type,attr,omitempty"`
	NotOnOrAfter string `xml:",attr,omitempty"`
	KeyInfo      *KeyInfo
}

type SubjectConfirmation struct {
	XMLName                 xml.Name
	Method                  string `xml:",attr"`
	NameID                  *NameID
	SubjectConfirmationData SubjectConfirmationData
}

type Condition struct {
	Type string `xml:"xsi:type,attr,omitempty"`
}

func (c *Condition) GetCondition() *Condition {
	return c
}

type BaseCondition interface {
	GetCondition() *Condition
}

func init() {
	types.Add("BaseCondition", reflect.TypeOf((*Condition)(nil)).Elem())
	types.Add("del:DelegationRestrictionType", reflect.TypeOf((*DelegateRestriction)(nil)).Elem())
	types.Add("rsa:RenewRestrictionType", reflect.TypeOf((*RenewRestriction)(nil)).Elem())
}

type Conditions struct {
	XMLName          xml.Name
	NotBefore        string            `xml:",attr"`
	NotOnOrAfter     string            `xml:",attr"`
	ProxyRestriction *ProxyRestriction `xml:",omitempty"`
	Condition        []BaseCondition   `xml:",omitempty"`
}

func (c *Conditions) C14N() string {
	names := []*xml.Name{
		&c.XMLName,
	}

	if c.ProxyRestriction != nil {
		names = append(names, &c.ProxyRestriction.XMLName)
	}

	for i := range c.Condition {
		switch r := c.Condition[i].(type) {
		case *DelegateRestriction:
			names = append(names, &r.XMLName, &r.Delegate.NameID.XMLName)
			r.NS = XSI
			r.Type = "del:DelegationRestrictionType"
			r.Delegate.NS = r.Delegate.XMLName.Space
			r.Delegate.XMLName = xml.Name{Local: "del:Delegate"}
		case *RenewRestriction:
			names = append(names, &r.XMLName)
			r.NS = XSI
			r.Type = "rsa:RenewRestrictionType"
		}
	}

	return mkns("saml2", c, names...)
}

type ProxyRestriction struct {
	XMLName xml.Name
	Count   int32 `xml:",attr"`
}

type RenewRestriction struct {
	XMLName xml.Name
	NS      string `xml:"xmlns:xsi,attr,omitempty"`
	Count   int32  `xml:",attr"`
	Condition
}

type Delegate struct {
	XMLName           xml.Name
	NS                string `xml:"xmlns:del,attr,omitempty"`
	DelegationInstant string `xml:",attr"`
	NameID            NameID
}

type DelegateRestriction struct {
	XMLName xml.Name
	NS      string `xml:"xmlns:xsi,attr,omitempty"`
	Condition
	Delegate Delegate
}

type AuthnStatement struct {
	XMLName      xml.Name
	AuthnInstant string `xml:",attr"`
	AuthnContext struct {
		XMLName              xml.Name
		AuthnContextClassRef struct {
			XMLName xml.Name
			Value   string `xml:",innerxml"`
		}
	}
}

func (a *AuthnStatement) C14N() string {
	return mkns("saml2", a, &a.XMLName, &a.AuthnContext.XMLName, &a.AuthnContext.AuthnContextClassRef.XMLName)
}

type AttributeStatement struct {
	XMLName   xml.Name
	Attribute []Attribute
}

func (a *AttributeStatement) C14N() string {
	c14n := []string{"<saml2:AttributeStatement>"}
	for i := range a.Attribute {
		c14n = append(c14n, a.Attribute[i].C14N())
	}
	c14n = append(c14n, "</saml2:AttributeStatement>")
	return strings.Join(c14n, "")
}

type AttributeValue struct {
	XMLName xml.Name
	Type    string `xml:"type,attr,typeattr"`
	Value   string `xml:",innerxml"`
}

func (a *AttributeValue) C14N() string {
	return fmt.Sprintf(`<saml2:AttributeValue xmlns:xsi="%s" xsi:type="%s">%s</saml2:AttributeValue>`, XSI, a.Type, a.Value)
}

type Attribute struct {
	XMLName        xml.Name
	FriendlyName   string `xml:",attr"`
	Name           string `xml:",attr"`
	NameFormat     string `xml:",attr"`
	AttributeValue []AttributeValue
}

func (a *Attribute) C14N() string {
	c14n := []string{
		fmt.Sprintf(`<saml2:Attribute FriendlyName="%s" Name="%s" NameFormat="%s">`, a.FriendlyName, a.Name, a.NameFormat),
	}

	for i := range a.AttributeValue {
		c14n = append(c14n, a.AttributeValue[i].C14N())
	}

	c14n = append(c14n, `</saml2:Attribute>`)

	return strings.Join(c14n, "")
}

type Lifetime struct {
	Created string `xml:"http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd Created"`
	Expires string `xml:"http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd Expires"`
}

func (t *Lifetime) C14N() string {
	return fmt.Sprintf(`<Lifetime><wsu:Created>%s</wsu:Created><wsu:Expires>%s</wsu:Expires></Lifetime>`, t.Created, t.Expires)
}

type Renewing struct {
	Allow bool `xml:",attr"`
	OK    bool `xml:",attr"`
}

type UseKey struct {
	Sig string `xml:",attr"`
}

type Target struct {
	Token string `xml:",innerxml"`
}

type RequestSecurityToken struct {
	TokenType          string    `xml:",omitempty"`
	RequestType        string    `xml:",omitempty"`
	Lifetime           *Lifetime `xml:",omitempty"`
	Renewing           *Renewing `xml:",omitempty"`
	Delegatable        bool      `xml:",omitempty"`
	KeyType            string    `xml:",omitempty"`
	SignatureAlgorithm string    `xml:",omitempty"`
	UseKey             *UseKey   `xml:",omitempty"`
	ActAs              *Target   `xml:",omitempty"`
	ValidateTarget     *Target   `xml:",omitempty"`
	RenewTarget        *Target   `xml:",omitempty"`
}

func Unmarshal(data []byte, v any) error {
	dec := xml.NewDecoder(bytes.NewReader(data))
	dec.TypeFunc = types.TypeFunc()
	return dec.Decode(v)
}

// toString returns an XML encoded RequestSecurityToken.
// When c14n is true, returns the canonicalized ActAs.Assertion which is required to sign the Issue request.
// When c14n is false, returns the original content of the ActAs.Assertion.
// The original content must be used within the request Body, as it has its own signature.
func (r *RequestSecurityToken) toString(c14n bool) string {
	actas := ""
	if r.ActAs != nil {
		token := r.ActAs.Token
		if c14n {
			var a Assertion
			err := Unmarshal([]byte(r.ActAs.Token), &a)
			if err != nil {
				log.Printf("decode ActAs: %s", err)
			}
			token = a.C14N()
		}

		actas = fmt.Sprintf(`<wst:ActAs xmlns:wst="http://docs.oasis-open.org/ws-sx/ws-trust/200802">%s</wst:ActAs>`, token)
	}

	body := []string{
		fmt.Sprintf(`<RequestSecurityToken xmlns="http://docs.oasis-open.org/ws-sx/ws-trust/200512">`),
		fmt.Sprintf(`<TokenType>%s</TokenType>`, r.TokenType),
		fmt.Sprintf(`<RequestType>%s</RequestType>`, r.RequestType),
		r.Lifetime.C14N(),
	}

	if r.RenewTarget == nil {
		body = append(body,
			fmt.Sprintf(`<Renewing Allow="%t" OK="%t"></Renewing>`, r.Renewing.Allow, r.Renewing.OK),
			fmt.Sprintf(`<Delegatable>%t</Delegatable>`, r.Delegatable),
			actas,
			fmt.Sprintf(`<KeyType>%s</KeyType>`, r.KeyType),
			fmt.Sprintf(`<SignatureAlgorithm>%s</SignatureAlgorithm>`, r.SignatureAlgorithm),
			fmt.Sprintf(`<UseKey Sig="%s"></UseKey>`, r.UseKey.Sig))
	} else {
		token := r.RenewTarget.Token
		if c14n {
			var a Assertion
			err := Unmarshal([]byte(r.RenewTarget.Token), &a)
			if err != nil {
				log.Printf("decode Renew: %s", err)
			}
			token = a.C14N()
		}

		body = append(body,
			fmt.Sprintf(`<UseKey Sig="%s"></UseKey>`, r.UseKey.Sig),
			fmt.Sprintf(`<RenewTarget>%s</RenewTarget>`, token))
	}

	return strings.Join(append(body, `</RequestSecurityToken>`), "")
}

func (r *RequestSecurityToken) C14N() string {
	return r.toString(true)
}

func (r *RequestSecurityToken) String() string {
	return r.toString(false)
}

type RequestSecurityTokenResponseCollection struct {
	RequestSecurityTokenResponse RequestSecurityTokenResponse
}

type RequestSecurityTokenResponse struct {
	RequestedSecurityToken RequestedSecurityToken
	Lifetime               *Lifetime `xml:"http://docs.oasis-open.org/ws-sx/ws-trust/200512 Lifetime"`
}

type RequestedSecurityToken struct {
	Assertion string `xml:",innerxml"`
}

type RequestSecurityTokenBody struct {
	Req    *RequestSecurityToken                   `xml:"http://docs.oasis-open.org/ws-sx/ws-trust/200512 RequestSecurityToken,omitempty"`
	Res    *RequestSecurityTokenResponseCollection `xml:"http://docs.oasis-open.org/ws-sx/ws-trust/200512 RequestSecurityTokenResponseCollection,omitempty"`
	Fault_ *soap.Fault                             `xml:"http://schemas.xmlsoap.org/soap/envelope/ Fault,omitempty"`
}

func (b *RequestSecurityTokenBody) Fault() *soap.Fault { return b.Fault_ }

func (b *RequestSecurityTokenBody) RequestSecurityToken() *RequestSecurityToken { return b.Req }

func (r *RequestSecurityToken) Action() string {
	kind := path.Base(r.RequestType)
	return "http://docs.oasis-open.org/ws-sx/ws-trust/200512/RST/" + kind
}

func Issue(ctx context.Context, r soap.RoundTripper, req *RequestSecurityToken) (*RequestSecurityTokenResponseCollection, error) {
	var reqBody, resBody RequestSecurityTokenBody

	reqBody.Req = req

	if err := r.RoundTrip(ctx, &reqBody, &resBody); err != nil {
		return nil, err
	}

	return resBody.Res, nil
}

type RenewSecurityTokenBody struct {
	Req    *RequestSecurityToken         `xml:"http://docs.oasis-open.org/ws-sx/ws-trust/200512 RequestSecurityToken,omitempty"`
	Res    *RequestSecurityTokenResponse `xml:"http://docs.oasis-open.org/ws-sx/ws-trust/200512 RequestSecurityTokenResponse,omitempty"`
	Fault_ *soap.Fault                   `xml:"http://schemas.xmlsoap.org/soap/envelope/ Fault,omitempty"`
}

func (b *RenewSecurityTokenBody) Fault() *soap.Fault { return b.Fault_ }

func (b *RenewSecurityTokenBody) RequestSecurityToken() *RequestSecurityToken { return b.Req }

func Renew(ctx context.Context, r soap.RoundTripper, req *RequestSecurityToken) (*RequestSecurityTokenResponse, error) {
	var reqBody, resBody RenewSecurityTokenBody

	reqBody.Req = req

	if err := r.RoundTrip(ctx, &reqBody, &resBody); err != nil {
		return nil, err
	}

	return resBody.Res, nil
}

// Marshal panics if xml.Marshal returns an error
func Marshal(val any) string {
	b, err := xml.Marshal(val)
	if err != nil {
		panic(err)
	}
	return string(b)
}

// mkns prepends the given namespace to xml.Name.Local and returns obj encoded as xml.
// Note that the namespace is required when encoding, but the namespace prefix must not be
// present when decoding as Go's decoding does not handle namespace prefix.
func mkns(ns string, obj any, name ...*xml.Name) string {
	ns += ":"
	for i := range name {
		name[i].Space = ""
		if !strings.HasPrefix(name[i].Local, ns) {
			name[i].Local = ns + name[i].Local
		}
	}

	return Marshal(obj)
}
