// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cert

import (
        "crypto"
        "crypto/rand"
        "crypto/x509"
        "crypto/x509/pkix"
        "encoding/asn1"
        "encoding/json"
        "encoding/pem"
        "fmt"
        "math/big"
        "time"

        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken"
)

// CreateX509Cert generates a self-signed x509 cert from a PK token
//   - OP 'sub' claim is mapped to the CN and SANs fields
//   - User public key is mapped to the RawSubjectPublicKeyInfo field
//   - Raw PK token is mapped to the SubjectKeyId field
func CreateX509Cert(pkToken *pktoken.PKToken, signer crypto.Signer) ([]byte, error) {
        template, err := PktToX509Template(pkToken)
        if err != nil {
                return nil, fmt.Errorf("error creating X.509 template: %w", err)
        }

        // create a self-signed X.509 certificate
        certDER, err := x509.CreateCertificate(rand.Reader, template, template, signer.Public(), signer)
        if err != nil {
                return nil, fmt.Errorf("error creating X.509 certificate: %w", err)
        }
        certBlock := &pem.Block{Type: "CERTIFICATE", Bytes: certDER}
        return pem.EncodeToMemory(certBlock), nil
}

// PktToX509Template takes a PK Token and returns a X.509 certificate template
// with the fields of the template set to the values in the X509
func PktToX509Template(pkt *pktoken.PKToken) (*x509.Certificate, error) {
        pktJson, err := json.Marshal(pkt)
        if err != nil {
                return nil, fmt.Errorf("error marshalling PK token to JSON: %w", err)
        }

        // get subject identifier from pk token
        idtClaims := new(oidc.OidcClaims)
        if err := json.Unmarshal(pkt.Payload, idtClaims); err != nil {
                return nil, err
        }

        cic, err := pkt.GetCicValues()
        if err != nil {
                return nil, err
        }
        upk := cic.PublicKey()
        var rawkey interface{} // This is the raw key, like *rsa.PrivateKey or *ecdsa.PrivateKey
        if err := upk.Raw(&rawkey); err != nil {
                return nil, err
        }

        // encode ephemeral public key
        ecPub, err := x509.MarshalPKIXPublicKey(rawkey)
        if err != nil {
                return nil, fmt.Errorf("error marshalling public key: %w", err)
        }
        template := &x509.Certificate{
                SerialNumber:            big.NewInt(1),
                Subject:                 pkix.Name{CommonName: idtClaims.Subject},
                RawSubjectPublicKeyInfo: ecPub,
                NotBefore:               time.Now(),
                NotAfter:                time.Now().Add(365 * 24 * time.Hour), // valid for 1 year
                KeyUsage:                x509.KeyUsageDigitalSignature,
                ExtKeyUsage:             []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning},
                BasicConstraintsValid:   true,
                DNSNames:                []string{idtClaims.Subject},
                IsCA:                    false,
                ExtraExtensions: []pkix.Extension{{
                        // OID for OIDC Issuer extension
                        Id:       asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 57264, 1, 1},
                        Critical: false,
                        Value:    []byte(idtClaims.Issuer),
                }},
                SubjectKeyId: pktJson,
        }

        return template, nil
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package choosers

import (
        "context"
        "embed"
        "errors"
        "fmt"
        "html/template"
        "io/fs"
        "net"
        "net/http"
        "net/http/httptest"
        "sort"
        "strings"

        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/util"
        "github.com/sirupsen/logrus"
)

//go:embed static/*
var staticFiles embed.FS

//go:embed chooser.tmpl
var chooserTemplateFile string

// To add support for an OP to the the WebChooser:
// 1. Add the OP to IssuerToName func
// 2. Add the OP to the html template file: `chooser.tmpl`
// 3. Add the OP to the data which is supplied to `chooserTemplate.Execute(w, data)`
//
// Note that the web chooser can only support BrowserOpenIdProvider

// TODO: This should be an enum that can also autogenerate what gets passed to the template

type WebChooser struct {
        OpList        []providers.BrowserOpenIdProvider
        opSelected    providers.BrowserOpenIdProvider
        OpenBrowser   bool
        useMockServer bool
        mockServer    *httptest.Server
        server        *http.Server
}

func NewWebChooser(opList []providers.BrowserOpenIdProvider, openBrowser bool) *WebChooser {
        return &WebChooser{
                OpList:        opList,
                OpenBrowser:   openBrowser,
                useMockServer: false,
        }
}

func (wc *WebChooser) ChooseOp(ctx context.Context) (providers.OpenIdProvider, error) {
        if wc.opSelected != nil {
                return nil, fmt.Errorf("provider has already been chosen")
        }

        providerMap := map[string]providers.BrowserOpenIdProvider{}
        for _, provider := range wc.OpList {
                if providerName, err := IssuerToName(provider.Issuer()); err != nil {
                        return nil, err
                } else {
                        if _, ok := providerMap[providerName]; ok {
                                return nil, fmt.Errorf("provider in web chooser found with duplicate issuer: %s", provider.Issuer())
                        }
                        providerMap[providerName] = provider
                }
        }

        opCh := make(chan providers.BrowserOpenIdProvider, 1)
        errCh := make(chan error, 1)

        mux := http.NewServeMux()
        staticContent, err := fs.Sub(staticFiles, "static")
        if err != nil {
                return nil, err
        }

        chooserTemplate, err := template.New("chooser-page").Parse(chooserTemplateFile)
        if err != nil {
                return nil, err
        }

        mux.HandleFunc("/chooser", func(w http.ResponseWriter, r *http.Request) {
                type Provider struct {
                        Name   string
                        Button string
                }

                data := struct {
                        Providers []Provider
                }{}

                sortedProviderNames := make([]string, 0)
                for providerName := range providerMap {
                        sortedProviderNames = append(sortedProviderNames, providerName)
                }
                // Sort the provider names
                sort.Strings(sortedProviderNames)

                for _, providerName := range sortedProviderNames {
                        if providerName == "google" {
                                data.Providers = append(data.Providers, Provider{
                                        Name:   providerName,
                                        Button: "google-light.svg",
                                })
                                continue
                        }
                        if providerName == "azure" {
                                data.Providers = append(data.Providers, Provider{
                                        Name:   providerName,
                                        Button: "azure-dark.svg",
                                })
                                continue
                        }
                        if providerName == "gitlab" {
                                data.Providers = append(data.Providers, Provider{
                                        Name:   providerName,
                                        Button: "gitlab-light.svg",
                                })
                                continue
                        }
                        if providerName == "hello" {
                                data.Providers = append(data.Providers, Provider{
                                        Name:   providerName,
                                        Button: "hello-dark.png",
                                })
                                continue
                        }
                        data.Providers = append(data.Providers, Provider{
                                Name:   providerName,
                                Button: "",
                        })

                }

                w.Header().Set("Content-Type", "text/html")
                if err := chooserTemplate.Execute(w, data); err != nil {
                        http.Error(w, err.Error(), http.StatusInternalServerError)
                }
        })
        mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.FS(staticContent))))
        mux.HandleFunc("/select/", func(w http.ResponseWriter, r *http.Request) {
                // Once we redirect to the OP localhost webserver, we can shutdown the web chooser localhost server
                shutdownServer := func() {
                        go func() { // Put this in a go func so that it will not block the redirect
                                if wc.server != nil {
                                        if err := wc.server.Shutdown(ctx); err != nil {
                                                logrus.Errorf("Failed to shutdown http server: %v", err)
                                        }
                                }
                        }()
                }
                defer shutdownServer()

                opName := r.URL.Query().Get("op")
                if opName == "" {
                        errorString := "missing op parameter"
                        http.Error(w, errorString, http.StatusBadRequest)
                        errCh <- errors.New(errorString)
                        return
                }
                if op, ok := providerMap[opName]; !ok {
                        errorString := fmt.Sprintf("unknown OpenID Provider: %s", opName)
                        http.Error(w, errorString, http.StatusBadRequest)
                        errCh <- errors.New(errorString)
                        return
                } else {
                        opCh <- op

                        redirectUriCh := make(chan string, 1)
                        op.ReuseBrowserWindowHook(redirectUriCh)

                        redirectUri := <-redirectUriCh
                        http.Redirect(w, r, redirectUri, http.StatusFound)
                }
        })

        if wc.useMockServer {
                wc.mockServer = httptest.NewUnstartedServer(mux)
                wc.mockServer.Start()
        } else {
                listener, err := net.Listen("tcp", "localhost:0")
                if err != nil {
                        return nil, fmt.Errorf("failed to bind to an available port: %w", err)
                }
                wc.server = &http.Server{Handler: mux}
                go func() {
                        err = wc.server.Serve(listener)
                        if err != nil && err != http.ErrServerClosed {
                                logrus.Error(err)
                        }
                }()

                var loginURI string
                if listener.Addr().(*net.TCPAddr).IP.String() == "127.0.0.1" {
                        // For consistency in output messages in our code base we use localhost rather than 127.0.0.1
                        port := listener.Addr().(*net.TCPAddr).Port
                        loginURI = fmt.Sprintf("http://localhost:%d/chooser", port)
                } else {
                        loginURI = fmt.Sprintf("http://%s/chooser", listener.Addr().String())
                }

                if wc.OpenBrowser {
                        logrus.Infof("Opening browser to %s", loginURI)
                        if err := util.OpenUrl(loginURI); err != nil {
                                logrus.Errorf("Failed to open url: %v", err)
                        }
                } else {
                        // If wc.OpenBrowser is false, tell the user what URL to open.
                        // This is useful when a user wants to use a different browser than the default one.
                        logrus.Infof("Open your browser to: %s ", loginURI)
                }
        }

        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case err := <-errCh:
                return nil, err
        case wc.opSelected = <-opCh:
                return wc.opSelected, nil
        }
}

func IssuerToName(issuer string) (string, error) {
        switch {
        case strings.HasPrefix(issuer, "https://accounts.google.com"):
                return "google", nil
        case strings.HasPrefix(issuer, "https://login.microsoftonline.com"):
                return "azure", nil
        case strings.HasPrefix(issuer, "https://gitlab.com"):
                return "gitlab", nil
        case strings.HasPrefix(issuer, "https://issuer.hello.coop"):
                return "hello", nil
        default:
                if strings.HasPrefix(issuer, "https://") {
                        // Returns issuer without the "https://" prefix and without any path remaining on the url
                        // e.g. https://accounts.google.com/fdsfa/fdsafsad -> accounts.google.com
                        return strings.Split(strings.TrimPrefix(issuer, "https://"), "/")[0], nil

                }
                return "", fmt.Errorf("invalid OpenID Provider issuer: %s", issuer)
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "context"
        "crypto"
        "fmt"
        "net/http"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"

        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/util"
        "github.com/openpubkey/openpubkey/verifier"
)

type OpenIdProvider = providers.OpenIdProvider

type BrowserOpenIdProvider = providers.BrowserOpenIdProvider

type PKTokenVerifier interface {
        VerifyPKToken(ctx context.Context, pkt *pktoken.PKToken, extraChecks ...verifier.Check) error
}

type OpkClient struct {
        Op           OpenIdProvider
        cosP         *CosignerProvider
        signer       crypto.Signer
        alg          jwa.KeyAlgorithm
        pkToken      *pktoken.PKToken
        refreshToken []byte
        accessToken  []byte
}

// ClientOpts contains options for constructing an OpkClient
type ClientOpts func(o *OpkClient)

// WithSigner allows the caller to inject their own signer and algorithm.
// Use this option if to generate to bring your own user key pair. If this
// option is not set the OpkClient constructor will automatically generate
// a signer, i.e., key pair.
// Example use:
//
//        signer, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
//        WithSigner(signer, jwa.ES256)
func WithSigner(signer crypto.Signer, alg jwa.KeyAlgorithm) ClientOpts {
        return func(o *OpkClient) {
                o.signer = signer
                o.alg = alg
        }
}

// WithCosignerProvider specifies what cosigner provider should be used to
// cosign the PK Token. If this is not specified then the cosigning setup
// is skipped.
func WithCosignerProvider(cosP *CosignerProvider) ClientOpts {
        return func(o *OpkClient) {
                o.cosP = cosP
        }
}

// New returns a new client.OpkClient. The op argument should be the
// OpenID Provider you want to authenticate against.
func New(op OpenIdProvider, opts ...ClientOpts) (*OpkClient, error) {
        client := &OpkClient{
                Op:     op,
                signer: nil,
                alg:    nil,
        }

        for _, applyOpt := range opts {
                applyOpt(client)
        }

        if client.alg == nil && client.signer != nil {
                return nil, fmt.Errorf("signer specified but alg is nil, must specify alg of signer")
        }

        if client.signer == nil {
                // Generate signer for specified alg. If no alg specified, defaults to ES256
                if client.alg == nil {
                        client.alg = jwa.ES256
                }

                signer, err := util.GenKeyPair(client.alg)
                if err != nil {
                        return nil, fmt.Errorf("failed to create key pair for client: %w ", err)
                }
                client.signer = signer
        }

        return client, nil
}

type AuthOptsStruct struct {
        extraClaims map[string]any
}
type AuthOpts func(a *AuthOptsStruct)

// WithExtraClaim specifies additional values to be included in the
// CIC. These claims will be include in the CIC protected header and
// will be hashed into the commitment claim in the ID Token. The
// commitment claim is typically the nonce or aud claim in the ID Token.
// Example use:
//
//        WithExtraClaim("claimKey", "claimValue")
func WithExtraClaim(k string, v string) AuthOpts {
        return func(a *AuthOptsStruct) {
                if a.extraClaims == nil {
                        a.extraClaims = map[string]any{}
                }
                a.extraClaims[k] = v
        }
}

// Auth returns a PK Token by running the OpenPubkey protocol. It will first
// authenticate to the configured OpenID Provider (OP) and receive an ID Token.
// Using this ID Token it will generate a PK Token. If a Cosigner has been
// configured it will also attempt to get the PK Token cosigned.
func (o *OpkClient) Auth(ctx context.Context, opts ...AuthOpts) (*pktoken.PKToken, error) {
        authOpts := &AuthOptsStruct{
                extraClaims: map[string]any{},
        }
        for _, applyOpt := range opts {
                applyOpt(authOpts)
        }

        // If no Cosigner is set then do standard OIDC authentication
        if o.cosP == nil {
                pkt, err := o.oidcAuth(ctx, o.signer, o.alg, authOpts.extraClaims)
                if err != nil {
                        return nil, err
                }
                o.pkToken = pkt
                return o.pkToken.DeepCopy()
        }

        // If a Cosigner is set then check that will support doing Cosigner auth
        if browserOp, ok := o.Op.(BrowserOpenIdProvider); !ok {
                return nil, fmt.Errorf("OP supplied does not have support for MFA Cosigner")
        } else {
                redirCh := make(chan string, 1)

                browserOp.HookHTTPSession(func(w http.ResponseWriter, r *http.Request) {
                        redirectUri := <-redirCh
                        http.Redirect(w, r, redirectUri, http.StatusFound)
                })

                pkt, err := o.oidcAuth(ctx, o.signer, o.alg, authOpts.extraClaims)
                if err != nil {
                        return nil, err
                }
                pktCos, err := o.cosP.RequestToken(ctx, o.signer, pkt, redirCh)
                if err != nil {
                        return nil, err
                }
                o.pkToken = pktCos
                return o.pkToken.DeepCopy()
        }
}

// oidcAuth performs the OpenIdConnect part of the protocol.
// Auth is the exposed function that should be called.
func (o *OpkClient) oidcAuth(
        ctx context.Context,
        signer crypto.Signer,
        alg jwa.KeyAlgorithm,
        extraClaims map[string]any,
) (*pktoken.PKToken, error) {
        // keep track of any additional verifierChecks for the verifier
        verifierChecks := []verifier.Check{}

        // Use our signing key to generate a JWK key and set the "alg" header
        jwkKey, err := jwk.PublicKeyOf(signer.Public())
        if err != nil {
                return nil, err
        }
        err = jwkKey.Set(jwk.AlgorithmKey, alg)
        if err != nil {
                return nil, err
        }

        // Use provided public key to generate client instance claims
        cic, err := clientinstance.NewClaims(jwkKey, extraClaims)
        if err != nil {
                return nil, fmt.Errorf("failed to instantiate client instance claims: %w", err)
        }

        tokens, err := o.Op.RequestTokens(ctx, cic)
        if err != nil {
                return nil, fmt.Errorf("error requesting OIDC tokens from OpenID Provider: %w", err)
        }
        idToken := tokens.IDToken
        o.refreshToken = tokens.RefreshToken
        o.accessToken = tokens.AccessToken

        // Sign over the payload from the ID token and client instance claims
        cicToken, err := cic.Sign(signer, alg, idToken)
        if err != nil {
                return nil, fmt.Errorf("error creating cic token: %w", err)
        }

        // Combine our ID token and signature over the cic to create our PK Token
        pkt, err := pktoken.New(idToken, cicToken)
        if err != nil {
                return nil, fmt.Errorf("error creating PK Token: %w", err)
        }

        pktVerifier, err := verifier.New(o.Op)
        if err != nil {
                return nil, err
        }

        if err := pktVerifier.VerifyPKToken(ctx, pkt, verifierChecks...); err != nil {
                return nil, fmt.Errorf("error verifying PK Token: %w", err)
        }
        return pkt, nil
}

// Refresh uses a Refresh Token to request a fresh ID Token and Access Token from an OpenID Provider.
// It provides a way to refresh the Access and ID Tokens for an OpenID Provider that supports refresh requests,
// allowing the client to continue making authenticated requests without requiring the user to re-authenticate.
func (o *OpkClient) Refresh(ctx context.Context) (*pktoken.PKToken, error) {
        if tokensOp, ok := o.Op.(providers.RefreshableOpenIdProvider); ok {
                if o.refreshToken == nil {
                        return nil, fmt.Errorf("no refresh token set")
                }
                if o.pkToken == nil {
                        return nil, fmt.Errorf("no PK Token set, run Auth() to create a PK Token first")
                }
                tokens, err := tokensOp.RefreshTokens(ctx, o.refreshToken)
                if err != nil {
                        return nil, fmt.Errorf("error requesting ID token: %w", err)
                }
                o.pkToken.FreshIDToken = tokens.IDToken
                o.refreshToken = tokens.RefreshToken
                o.accessToken = tokens.AccessToken

                return o.pkToken.DeepCopy()
        }
        return nil, fmt.Errorf("OP (issuer=%s) does not support OIDC refresh requests", o.Op.Issuer())
}

// GetOp returns the OpenID Provider the OpkClient has been configured to use
func (o *OpkClient) GetOp() OpenIdProvider {
        return o.Op
}

// GetCosP returns the MFA Cosigner Provider the OpkClient has been
// configured to use
func (o *OpkClient) GetCosP() *CosignerProvider {
        return o.cosP
}

// GetSigner returns the client's key pair (Public Key, Signing Key)
func (o *OpkClient) GetSigner() crypto.Signer {
        return o.signer
}

// GetAlg returns the algorithm of the client's key pair
// (Public Key, Signing Key)
func (o *OpkClient) GetAlg() jwa.KeyAlgorithm {
        return o.alg
}

func (o *OpkClient) SetPKToken(pkt *pktoken.PKToken) {
        o.pkToken = pkt
}

// GetPKToken returns a deep copy of client's current PK Token
func (o *OpkClient) GetPKToken() (*pktoken.PKToken, error) {
        return o.pkToken.DeepCopy()
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "context"
        "crypto"
        "crypto/rand"
        "encoding/hex"
        "encoding/json"
        "fmt"
        "io"
        "net"
        "net/http"
        "net/url"
        "strings"
        "time"

        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/cosigner/msgs"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/util"
        "github.com/sirupsen/logrus"
)

type CosignerProvider struct {
        Issuer       string
        CallbackPath string
}

func (c *CosignerProvider) RequestToken(ctx context.Context, signer crypto.Signer, pkt *pktoken.PKToken, redirCh chan string) (*pktoken.PKToken, error) {
        // Find an unused port
        listener, err := net.Listen("tcp", ":0")
        if err != nil {
                return nil, fmt.Errorf("failed to bind to an available port: %w", err)
        }

        port := listener.Addr().(*net.TCPAddr).Port
        host := fmt.Sprintf("localhost:%d", port)
        redirectURI := fmt.Sprintf("http://%s%s", host, c.CallbackPath)

        // We set the buffer size to one and then in the CallbackPath handler we
        // ensure only said either 0 or 1 message to a channel before returning.
        // This prevents blocking inside CallbackPath handler when it attempts to
        // write to the channel. If the callbackPath handler is called twice by the
        // user's web browser the second call will block on a channel until the cxt
        // is marked as done.
        sigCh := make(chan []byte, 1)
        errCh := make(chan error, 1)

        // This is where we get the authcode from the Cosigner
        mux := http.NewServeMux()
        mux.Handle(c.CallbackPath,
                http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        cosSig, err := func() ([]byte, error) {
                                // Get authcode from Cosigner via Cosigner redirecting user's browser window
                                params := r.URL.Query()
                                if _, ok := params["authcode"]; !ok {
                                        return nil, fmt.Errorf("cosigner did not return an authcode in the URI")
                                }
                                authcode := params["authcode"][0] // This is the authcode issued by the cosigner not the OP

                                // Sign authcode from cosigner under PK Token and send signed authcode to Cosigner
                                sig2, err := pkt.NewSignedMessage([]byte(authcode), signer)
                                if err != nil {
                                        return nil, fmt.Errorf("cosigner client hit error when building authcode URI: %w", err)
                                }
                                authcodeSigUri, err := c.authcodeURI(sig2)
                                if err != nil {
                                        return nil, fmt.Errorf("cosigner client hit error when building authcode URI: %w", err)
                                }
                                res, err := http.Get(authcodeSigUri)
                                if err != nil {
                                        return nil, fmt.Errorf("error requesting MFA cosigner signature: %w", err)
                                }

                                // Receive response from Cosigner that has cosigner signature on PK Token
                                resBody, err := io.ReadAll(res.Body)
                                if err != nil {
                                        return nil, fmt.Errorf("error reading MFA cosigner signature response: %w", err)
                                }
                                cosSig, err := util.Base64DecodeForJWT(resBody)
                                if err != nil {
                                        return nil, fmt.Errorf("error reading MFA cosigner signature response: %w", err)
                                }
                                // Success
                                return cosSig, nil
                        }()

                        if err != nil {
                                // Write the error message to the user
                                if _, err := w.Write([]byte(err.Error())); err != nil {
                                        logrus.Error(err)
                                }

                                select {
                                case errCh <- err:
                                case <-ctx.Done():
                                        return
                                }
                        } else {
                                if _, err := w.Write([]byte("You may now close this window")); err != nil {
                                        logrus.Error(err)
                                }

                                select {
                                case sigCh <- cosSig:
                                case <-ctx.Done():
                                        return
                                }
                        }

                }),
        )

        server := &http.Server{
                Addr:    host,
                Handler: mux,
        }
        logrus.Infof("listening on http://%s/", host)
        logrus.Info("press ctrl+c to stop")
        go func() {
                err := server.Serve(listener)
                if err != nil && err != http.ErrServerClosed {
                        logrus.Error(err)
                }
        }()
        defer func() {
                if err := server.Shutdown(ctx); err != nil {
                        logrus.Error(err)
                }
        }()

        pktJson, err := json.Marshal(pkt)
        if err != nil {
                return nil, fmt.Errorf("cosigner client hit error serializing PK Token: %w", err)
        }

        initAuthMsgJson, nonce, err := c.CreateInitAuthSig(redirectURI)
        if err != nil {
                return nil, fmt.Errorf("hit error creating init auth signed message: %w", err)
        }
        sig1, err := pkt.NewSignedMessage(initAuthMsgJson, signer)
        if err != nil {
                return nil, fmt.Errorf("cosigner client hit error init auth signed message: %w", err)
        }

        redirUri, err := c.initAuthURI(pktJson, sig1)
        if err != nil {
                return nil, fmt.Errorf("cosigner client hit error when building init auth URI: %w", err)
        }

        select {
        // Trigger redirect of user's browser window to a URI controlled by the Cosigner sending the PK Token in the URI
        case redirCh <- redirUri:
        case <-ctx.Done():
                return nil, ctx.Err()
        }

        select {
        case cosSig := <-sigCh: // Received cosigner signature
                // To be safe we perform these checks before adding the cosSig to the pktoken
                if err := c.ValidateCos(cosSig, nonce, redirectURI); err != nil {
                        return nil, err
                }
                if err := pkt.AddSignature(cosSig, pktoken.COS); err != nil {
                        return nil, fmt.Errorf("error in adding cosigner signature to PK Token: %w", err)
                }
                return pkt, nil
        case err := <-errCh:
                return nil, err
        case <-ctx.Done():
                return nil, ctx.Err()
        }
}

func (c *CosignerProvider) initAuthURI(pktJson []byte, sig1 []byte) (string, error) {
        pktB63 := util.Base64EncodeForJWT(pktJson)
        if uri, err := url.Parse(c.Issuer); err != nil {
                return "", err
        } else {
                uri := uri.JoinPath("mfa-auth-init")
                v := uri.Query()
                v.Add("pkt", string(pktB63))
                v.Add("sig1", string(sig1))
                uri.RawQuery = v.Encode()

                // URI Should be: https://<issuer>/mfa-auth-init?pkt=<pktJsonB64>&sig1=<sig1>
                return uri.String(), nil
        }
}

func (c *CosignerProvider) authcodeURI(sig2 []byte) (string, error) {
        if uri, err := url.Parse(c.Issuer); err != nil {
                return "", err
        } else {
                uri := uri.JoinPath("sign")
                v := uri.Query()
                v.Add("sig2", string(sig2))
                uri.RawQuery = v.Encode()

                // URI Should be: https://<issuer>/sign?&sig2=<sig2>
                return uri.String(), nil
        }

}

func (c *CosignerProvider) ValidateCos(cosSig []byte, expectedNonce string, expectedRedirectURI string) error {
        cosSigParsed, err := jws.Parse(cosSig)
        if err != nil {
                return fmt.Errorf("failed to parse Cosigner signature: %w", err)
        }

        if len(cosSigParsed.Signatures()) != 1 {
                return fmt.Errorf("the Cosigner signature does not have the correct number of signatures: %w", err)
        }

        ph := cosSigParsed.Signatures()[0].ProtectedHeaders()
        nonceRet, ok := ph.Get("nonce")
        if !ok {
                return fmt.Errorf("nonce not set in Cosigner signature protected header")
        }

        if expectedNonce != nonceRet {
                return fmt.Errorf("incorrect nonce set in Cosigner signature")
        }

        ruriRet, ok := ph.Get("ruri")
        if !ok {
                return fmt.Errorf("ruri (redirect URI) not set in Cosigner signature protected header")
        }

        if expectedRedirectURI != ruriRet {
                return fmt.Errorf("unexpected ruri (redirect URI) set in Cosigner signature, got %s expected %s", ruriRet, expectedRedirectURI)
        }

        issRet, ok := ph.Get("iss")
        if !ok {
                return fmt.Errorf("iss (Cosigner Issuer) not set in Cosigner signature protected header")
        }

        if c.Issuer != issRet {
                return fmt.Errorf("unexpected iss (Cosigner Issuer) set in Cosigner signature, expected %s", c.Issuer)
        }
        return nil
}

// CreateInitAuthSig generates a random nonce, validates the redirectURI,
// creates an InitMFAAuth message, marshals it to JSON,
// and returns the JSON message along with the nonce.
func (c *CosignerProvider) CreateInitAuthSig(redirectURI string) ([]byte, string, error) {
        bits := 256
        rBytes := make([]byte, bits/8)
        if _, err := rand.Read(rBytes); err != nil {
                return nil, "", err
        }
        if !strings.HasSuffix(redirectURI, c.CallbackPath) {
                return nil, "", fmt.Errorf("redirectURI (%s) does not end in expected callbackPath (%s)", redirectURI, c.CallbackPath)
        }

        nonce := hex.EncodeToString(rBytes)

        msg := msgs.InitMFAAuth{
                Issuer:      c.Issuer,
                RedirectUri: redirectURI,
                TimeSigned:  time.Now().Unix(),
                Nonce:       nonce,
        }
        msgJson, err := json.Marshal(msg)
        if err != nil {
                return nil, "", err
        }
        return msgJson, nonce, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cosigner

import (
        "crypto"
        "encoding/json"
        "fmt"
        "time"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/cosigner/msgs"
        "github.com/openpubkey/openpubkey/pktoken"
)

type AuthCosigner struct {
        Cosigner
        Issuer         string
        KeyID          string
        AuthStateStore AuthStateStore
}

func New(signer crypto.Signer, alg jwa.SignatureAlgorithm, issuer, keyID string, store AuthStateStore) (*AuthCosigner, error) {
        return &AuthCosigner{
                Cosigner: Cosigner{
                        Alg:    alg,
                        Signer: signer},
                Issuer:         issuer,
                KeyID:          keyID,
                AuthStateStore: store,
        }, nil
}

func (c *AuthCosigner) InitAuth(pkt *pktoken.PKToken, sig []byte) (string, error) {
        msg, err := pkt.VerifySignedMessage(sig)
        if err != nil {
                return "", fmt.Errorf("failed to verify sig: %w", err)
        }
        var initMFAAuth *msgs.InitMFAAuth
        if err := json.Unmarshal(msg, &initMFAAuth); err != nil {
                return "", fmt.Errorf("failed to parse InitMFAAuth message: %w", err)
        } else if initMFAAuth.Issuer != c.Issuer {
                return "", fmt.Errorf("signed message is for wrong cosigner, got issuer=(%s), expected issuer=(%s)", initMFAAuth.Issuer, c.Issuer)
        } else if time.Since(time.Unix(initMFAAuth.TimeSigned, 0)).Minutes() > 2 {
                return "", fmt.Errorf("timestamp (%d) in InitMFAAuth message too old, current time is (%d)", initMFAAuth.TimeSigned, time.Now().Unix())
        } else if time.Until(time.Unix(initMFAAuth.TimeSigned, 0)).Minutes() > 2 {
                return "", fmt.Errorf("timestamp (%d) in InitMFAAuth message too far in the future, current time is (%d)", initMFAAuth.TimeSigned, time.Now().Unix())
        } else if authID, err := c.AuthStateStore.CreateNewAuthSession(pkt, initMFAAuth.RedirectUri, initMFAAuth.Nonce); err != nil {
                return "", err
        } else {
                return authID, nil
        }
}

func (c *AuthCosigner) NewAuthcode(authID string) (string, error) {
        return c.AuthStateStore.CreateAuthcode(authID)
}

func (c *AuthCosigner) RedeemAuthcode(sig []byte) ([]byte, error) {
        msg, err := jws.Parse(sig)
        if err != nil {
                return nil, fmt.Errorf("failed to parse sig: %s", err)
        }
        authcode := string(msg.Payload())

        // We need redemption to be inside of our mutexes to ensure the same authcode can't be redeemed if requested at the same moment
        if authState, authID, err := c.AuthStateStore.RedeemAuthcode(authcode); err != nil {
                return nil, err
        } else {
                pkt := authState.Pkt
                _, err := pkt.VerifySignedMessage(sig) // We check this after redeeming the authcode, so can't try the same correct authcode twice
                if err != nil {
                        return nil, fmt.Errorf("error verifying sig: %w", err)
                }
                return c.IssueSignature(pkt, authState, authID)
        }
}

func (c *AuthCosigner) IssueSignature(pkt *pktoken.PKToken, authState AuthState, authID string) ([]byte, error) {

        protected := pktoken.CosignerClaims{
                Issuer:      c.Issuer,
                KeyID:       c.KeyID,
                Algorithm:   c.Alg.String(),
                AuthID:      authID,
                AuthTime:    time.Now().Unix(),
                IssuedAt:    time.Now().Unix(),
                Expiration:  time.Now().Add(time.Hour).Unix(),
                RedirectURI: authState.RedirectURI,
                Nonce:       authState.Nonce,
                Typ:         string(pktoken.COS),
        }

        // Now that our mfa has authenticated the user, we can add our signature
        return c.Cosign(pkt, protected)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cosigner

import (
        "crypto"
        "crypto/hmac"
        "encoding/binary"
        "encoding/hex"
        "fmt"
        "sync/atomic"
)

type AuthIDIssuer struct {
        authIdIter atomic.Uint64
        hmacKey    []byte
}

func NewAuthIDIssuer(hmacKey []byte) *AuthIDIssuer {
        return &AuthIDIssuer{
                authIdIter: atomic.Uint64{},
                hmacKey:    hmacKey,
        }
}

func (i *AuthIDIssuer) CreateAuthID(timeNow uint64) (string, error) {
        authIdInt := i.authIdIter.Add(1)
        iterAndTime := []byte{}
        iterAndTime = binary.LittleEndian.AppendUint64(iterAndTime, uint64(authIdInt))
        iterAndTime = binary.LittleEndian.AppendUint64(iterAndTime, timeNow)
        mac := hmac.New(crypto.SHA3_256.New, i.hmacKey)
        if n, err := mac.Write(iterAndTime); err != nil {
                return "", err
        } else if n != 16 {
                return "", fmt.Errorf("unexpected number of bytes read by HMAC, expected 16, got %d", n)
        } else {
                return hex.EncodeToString(mac.Sum(nil)), nil
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cosigner

import (
        "encoding/json"
        "fmt"
        "strings"

        "github.com/openpubkey/openpubkey/pktoken"
)

type AuthState struct {
        Pkt              *pktoken.PKToken
        Issuer           string // ID Token issuer (iss)
        Aud              string // ID Token audience (aud)
        Sub              string // ID Token subject ID (sub)
        Username         string // ID Token email or username
        DisplayName      string // ID Token display name (or username if none given)
        RedirectURI      string // Redirect URI
        Nonce            string // Nonce supplied by user
        AuthcodeIssued   bool   // Has an authcode been issued for this auth session
        AuthcodeRedeemed bool   // Was the pkt cosigned
}

func NewAuthState(pkt *pktoken.PKToken, ruri string, nonce string) (*AuthState, error) {
        var claims struct {
                Issuer string `json:"iss"`
                Aud    any    `json:"aud"`
                Sub    string `json:"sub"`
                Email  string `json:"email"`
        }
        if err := json.Unmarshal(pkt.Payload, &claims); err != nil {
                return nil, fmt.Errorf("failed to unmarshal PK Token: %w", err)
        }
        // An audience can be a string or an array of strings.
        //
        // RFC-7519 JSON Web Token (JWT) says:
        // "In the general case, the "aud" value is an array of case-
        // sensitive strings, each containing a StringOrURI value.  In the
        // special case when the JWT has one audience, the "aud" value MAY be a
        // single case-sensitive string containing a StringOrURI value."
        // https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3
        var audience string
        switch t := claims.Aud.(type) {
        case string:
                audience = t
        case []any:
                audList := []string{}
                for _, v := range t {
                        audList = append(audList, v.(string))
                }
                audience = strings.Join(audList, ",")
        default:
                return nil, fmt.Errorf("failed to deserialize aud (audience) claim in ID Token: %T", t)
        }

        return &AuthState{
                Pkt:              pkt,
                Issuer:           claims.Issuer,
                Aud:              audience,
                Sub:              claims.Sub,
                Username:         claims.Email,
                DisplayName:      strings.Split(claims.Email, "@")[0], //TODO: Use full name from ID Token
                RedirectURI:      ruri,
                Nonce:            nonce,
                AuthcodeRedeemed: false,
                AuthcodeIssued:   false,
        }, nil

}

type UserKey struct {
        Issuer string // ID Token issuer (iss)
        Aud    string // ID Token audience (aud)
        Sub    string // ID Token subject ID (sub)
}

func (as AuthState) UserKey() UserKey {
        return UserKey{Issuer: as.Issuer, Aud: as.Aud, Sub: as.Sub}
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cosigner

import (
        "crypto"
        "encoding/json"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/openpubkey/openpubkey/pktoken"
)

type Cosigner struct {
        Alg    jwa.KeyAlgorithm
        Signer crypto.Signer
}

func (c *Cosigner) Cosign(pkt *pktoken.PKToken, cosClaims pktoken.CosignerClaims) ([]byte, error) {
        jsonBytes, err := json.Marshal(cosClaims)
        if err != nil {
                return nil, err
        }
        var headers map[string]any
        if err := json.Unmarshal(jsonBytes, &headers); err != nil {
                return nil, err
        }
        return pkt.SignToken(c.Signer, c.Alg, headers)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package cosigner

import (
        "context"
        "fmt"
        "time"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/discover"
        "github.com/openpubkey/openpubkey/pktoken"
)

type DefaultCosignerVerifier struct {
        issuer  string
        options CosignerVerifierOpts
}

type CosignerVerifierOpts struct {
        // Strict specifies whether or not a pk token MUST contain a signature by this cosigner.
        // Defaults to true.
        Strict *bool
        // Allows users to set custom function for discovering public key of Cosigner
        DiscoverPublicKey *discover.PublicKeyFinder
}

func NewCosignerVerifier(issuer string, options CosignerVerifierOpts) *DefaultCosignerVerifier {
        v := &DefaultCosignerVerifier{
                issuer:  issuer,
                options: options,
        }

        // If no custom DiscoverPublicKey function is set, set default
        if v.options.DiscoverPublicKey == nil {
                v.options.DiscoverPublicKey = discover.DefaultPubkeyFinder()
        }

        // If strict is not set, then default it to true
        if v.options.Strict == nil {
                v.options.Strict = new(bool)
                *v.options.Strict = true
        }

        return v
}

func (v *DefaultCosignerVerifier) Issuer() string {
        return v.issuer
}

func (v *DefaultCosignerVerifier) Strict() bool {
        return *v.options.Strict
}

func (v *DefaultCosignerVerifier) VerifyCosigner(ctx context.Context, pkt *pktoken.PKToken) error {
        if pkt.Cos == nil {
                return fmt.Errorf("no cosigner signature")
        }

        // Parse our header
        header, err := pkt.ParseCosignerClaims()
        if err != nil {
                return err
        }

        if v.issuer != header.Issuer {
                return fmt.Errorf("cosigner issuer (%s) doesn't match expected issuer (%s)", header.Issuer, v.issuer)
        }

        keyRecord, err := v.options.DiscoverPublicKey.ByKeyID(ctx, v.issuer, header.KeyID)
        if err != nil {
                return err
        }
        key := keyRecord.PublicKey
        alg := keyRecord.Alg

        // Check if it's expired
        if time.Now().After(time.Unix(header.Expiration, 0)) {
                return fmt.Errorf("cosigner signature expired")
        }
        if header.Algorithm != alg {
                return fmt.Errorf("key (kid=%s) has alg (%s) which doesn't match alg (%s) in protected", header.KeyID, alg, header.Algorithm)
        }
        jwsPubkey := jws.WithKey(jwa.KeyAlgorithmFrom(alg), key)
        _, err = jws.Verify(pkt.CosToken, jwsPubkey)

        return err
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mocks

import (
        "crypto/rand"
        "encoding/hex"
        "encoding/json"
        "fmt"
        "strings"
        "sync"
        "time"

        "github.com/openpubkey/openpubkey/cosigner"
        "github.com/openpubkey/openpubkey/pktoken"
)

// This is intended for testing purposes. The locking strategy used is not
// particularly efficient. Anyone building a Cosigner should use the interface
// above to replace this in-memory store with a database.
type AuthStateInMemoryStore struct {
        AuthIDIssuer     *cosigner.AuthIDIssuer
        AuthStateMap     map[string]*cosigner.AuthState
        AuthCodeMap      map[string]string
        AuthStateMapLock sync.RWMutex
        AuthcodeMapLock  sync.RWMutex
}

func NewAuthStateInMemoryStore(hmacKey []byte) *AuthStateInMemoryStore {
        return &AuthStateInMemoryStore{
                AuthStateMap:     make(map[string]*cosigner.AuthState),
                AuthCodeMap:      make(map[string]string),
                AuthcodeMapLock:  sync.RWMutex{},
                AuthStateMapLock: sync.RWMutex{},
                AuthIDIssuer:     cosigner.NewAuthIDIssuer(hmacKey),
        }
}

// Writes to the AuthState are not concurrency safe, do not write
func (s *AuthStateInMemoryStore) LookupAuthState(authID string) (*cosigner.AuthState, bool) {
        s.AuthStateMapLock.RLock()
        as, ok := s.AuthStateMap[authID]
        s.AuthStateMapLock.RUnlock()
        return as, ok // Pass by value to prevent writes to the original
}

func (s *AuthStateInMemoryStore) UpdateAuthState(authID string, authState cosigner.AuthState) error {
        s.AuthStateMapLock.Lock()
        defer s.AuthStateMapLock.Unlock()

        if _, ok := s.AuthStateMap[authID]; !ok {
                return fmt.Errorf("failed to upload auth session because authID specified matches no session")
        } else {
                s.AuthStateMap[authID] = &authState
                return nil
        }
}

func (s *AuthStateInMemoryStore) CreateNewAuthSession(pkt *pktoken.PKToken, ruri string, nonce string) (string, error) {
        var claims struct {
                Issuer string `json:"iss"`
                Aud    any    `json:"aud"`
                Sub    string `json:"sub"`
                Email  string `json:"email"`
        }
        if err := json.Unmarshal(pkt.Payload, &claims); err != nil {
                return "", fmt.Errorf("failed to unmarshal PK Token: %w", err)
        }
        // An audience can be a string or an array of strings.
        //
        // RFC-7519 JSON Web Token (JWT) says:
        // "In the general case, the "aud" value is an array of case-
        // sensitive strings, each containing a StringOrURI value.  In the
        // special case when the JWT has one audience, the "aud" value MAY be a
        // single case-sensitive string containing a StringOrURI value."
        // https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3
        var audience string
        switch t := claims.Aud.(type) {
        case string:
                audience = t
        case []any:
                audList := []string{}
                for _, v := range t {
                        audList = append(audList, v.(string))
                }
                audience = strings.Join(audList, ",")
        default:
                return "", fmt.Errorf("failed to deserialize aud (audience) claim in ID Token: %T", t)
        }

        authState := &cosigner.AuthState{
                Pkt:              pkt,
                Issuer:           claims.Issuer,
                Aud:              audience,
                Sub:              claims.Sub,
                Username:         claims.Email,
                DisplayName:      strings.Split(claims.Email, "@")[0], //TODO: Use full name from ID Token
                RedirectURI:      ruri,
                Nonce:            nonce,
                AuthcodeRedeemed: false,
                AuthcodeIssued:   false,
        }

        if authID, err := s.AuthIDIssuer.CreateAuthID(uint64(time.Now().Unix())); err != nil {
                return "", err
        } else {
                s.AuthStateMapLock.Lock()
                if _, ok := s.AuthStateMap[authID]; ok {
                        return "", fmt.Errorf("specified authID is already in use")
                }
                s.AuthStateMap[authID] = authState
                s.AuthStateMapLock.Unlock()
                return authID, nil
        }
}

func (s *AuthStateInMemoryStore) CreateAuthcode(authID string) (string, error) {
        authCodeBytes := make([]byte, 32)
        if _, err := rand.Read(authCodeBytes); err != nil {
                return "", err
        }
        authcode := hex.EncodeToString(authCodeBytes)

        // We take a full read write lock here to ensure we don't issue an authcode twice for the same session
        s.AuthStateMapLock.Lock()
        defer s.AuthStateMapLock.Unlock()

        if authState, ok := s.AuthStateMap[authID]; !ok {
                return "", fmt.Errorf("no such authID")
        } else if authState.AuthcodeIssued {
                return "", fmt.Errorf("authcode already issued for this authID")
        } else {
                s.AuthcodeMapLock.Lock()
                defer s.AuthcodeMapLock.Unlock()

                if _, ok := s.AuthCodeMap[authcode]; ok {
                        return "", fmt.Errorf("authcode collision implies randomness failure in RNG")
                }
                authState.AuthcodeIssued = true
                s.AuthCodeMap[authcode] = authID

                return authcode, nil
        }
}

func (s *AuthStateInMemoryStore) RedeemAuthcode(authcode string) (cosigner.AuthState, string, error) {
        s.AuthcodeMapLock.RLock()
        authID, authcodeFound := s.AuthCodeMap[authcode]
        s.AuthcodeMapLock.RUnlock()
        if !authcodeFound {
                return cosigner.AuthState{}, "", fmt.Errorf("invalid authcode")
        } else {
                s.AuthStateMapLock.Lock()
                defer s.AuthStateMapLock.Unlock()

                authState := s.AuthStateMap[authID]
                if !authState.AuthcodeIssued {
                        // This should never happen
                        return cosigner.AuthState{}, "", fmt.Errorf("no authcode issued for this authID")
                }
                if authState.AuthcodeRedeemed {
                        return cosigner.AuthState{}, "", fmt.Errorf("authcode has already been redeemed")
                }
                authState.AuthcodeRedeemed = true
                return *authState, authID, nil
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package discover

import (
        "context"
        "crypto"
        "crypto/ecdsa"
        "crypto/rsa"
        "encoding/json"
        "fmt"
        "io"
        "net/http"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/gq"
        "github.com/openpubkey/openpubkey/util"

        oidcclient "github.com/zitadel/oidc/v3/pkg/client"
)

type PublicKeyRecord struct {
        PublicKey crypto.PublicKey
        Alg       string
        Issuer    string
}

func NewPublicKeyRecord(key jwk.Key, issuer string) (*PublicKeyRecord, error) {
        var pubKey interface{}
        if key.Algorithm() == jwa.RS256 {
                pubKey = new(rsa.PublicKey)
        } else if key.Algorithm() == jwa.ES256 {
                pubKey = new(ecdsa.PublicKey)
        } else if key.Algorithm().String() == "" {
                // OPs such as azure (microsoft) do not specify alg in their JWKS. To
                // handle this case, assume no alg in JWKS means RSA as OIDC requires
                // OPs use RSA.
                pubKey = new(rsa.PublicKey)
        } else {
                return nil, fmt.Errorf("JWK has unsupported alg (%s)", key.Algorithm())
        }
        err := key.Raw(&pubKey)
        if err != nil {
                return nil, fmt.Errorf("failed to decode public key: %w", err)
        }

        var alg string
        if key.Algorithm().String() == "" {
                alg = jwa.RS256.String()
        } else {
                alg = key.Algorithm().String()
        }

        return &PublicKeyRecord{
                PublicKey: pubKey,
                Alg:       alg,
                Issuer:    issuer,
        }, nil
}

func DefaultPubkeyFinder() *PublicKeyFinder {
        return &PublicKeyFinder{
                JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                        return GetJwksByIssuer(ctx, issuer, nil)
                },
        }
}

type JwksFetchFunc func(ctx context.Context, issuer string) ([]byte, error)

type PublicKeyFinder struct {
        JwksFunc JwksFetchFunc
}

// GetJwksByIssuer fetches the JWKS from the issuer's JWKS endpoint found at the
// issuer's well-known configuration. It doesn't attempt to parse the response
// but instead returns the JSON bytes of the JWKS. If httpClient is nil, then
// http.DefaultClient is used when fetching.
func GetJwksByIssuer(ctx context.Context, issuer string, httpClient *http.Client) ([]byte, error) {
        if httpClient == nil {
                httpClient = http.DefaultClient
        }

        discConf, err := oidcclient.Discover(ctx, issuer, httpClient)
        if err != nil {
                return nil, fmt.Errorf("failed to call OIDC discovery endpoint: %w", err)
        }

        request, err := http.NewRequestWithContext(ctx, "GET", discConf.JwksURI, nil)
        if err != nil {
                return nil, err
        }

        response, err := httpClient.Do(request)
        if err != nil {
                return nil, err
        }
        defer response.Body.Close()

        resp, err := httpClient.Get(discConf.JwksURI)
        if err != nil {
                return nil, fmt.Errorf("failed to fetch to JWKS: %w", err)
        }
        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("received non-200 from JWKS URI: %s", http.StatusText(response.StatusCode))
        }
        return io.ReadAll(resp.Body)
}

func (f *PublicKeyFinder) fetchAndParseJwks(ctx context.Context, issuer string) (jwk.Set, error) {
        jwksJson, err := f.JwksFunc(ctx, issuer)
        if err != nil {
                return nil, fmt.Errorf(`failed to fetch JWKS: %w`, err)
        }
        jwks := jwk.NewSet()
        if err := json.Unmarshal(jwksJson, jwks); err != nil {
                return nil, fmt.Errorf(`failed to unmarshal JWKS: %w`, err)
        }
        return jwks, nil
}

// ByToken looks up an OP public key in the JWKS using the KeyID (kid) in the
// protected header from the supplied token.
func (f *PublicKeyFinder) ByToken(ctx context.Context, issuer string, token []byte) (*PublicKeyRecord, error) {
        jwt, err := jws.Parse(token)
        if err != nil {
                return nil, fmt.Errorf("error parsing JWK in JWKS: %w", err)
        }
        // a JWT is guaranteed to have exactly one signature
        headers := jwt.Signatures()[0].ProtectedHeaders()

        if headers.Algorithm() == gq.GQ256 {
                origHeadersJson, err := util.Base64DecodeForJWT([]byte(headers.KeyID()))
                if err != nil {
                        return nil, fmt.Errorf("error base64 decoding GQ kid: %w", err)
                }

                // If GQ then replace the GQ headers with the original headers
                err = json.Unmarshal(origHeadersJson, &headers)
                if err != nil {
                        return nil, fmt.Errorf("error unmarshalling GQ kid to original headers: %w", err)
                }
        }
        // Use the KeyID (kid) in the headers from the supplied token to look up the public key
        return f.ByKeyID(ctx, issuer, headers.KeyID())
}

// ByKeyID looks up an OP public key in the JWKS using the KeyID (kid) supplied.
// If no KeyID (kid) exists in the header and there is only one key in the JWKS,
// that key is returned. This  is useful for cases where an OP may not set a KeyID
// (kid) in the JWT header.
//
// The JWT RFC states that it is acceptable to not use a KeyID (kid) if there is
// only one key in the JWKS:
// "The "kid" (key ID) parameter is used to match a specific key.  This is used,
// for instance, to choose among a set of keys within a JWK Set
// during key rollover.  The structure of the "kid" value is
// unspecified.  When "kid" values are used within a JWK Set, different
// keys within the JWK Set SHOULD use distinct "kid" values.  (One
// example in which different keys might use the same "kid" value is if
// they have different "kty" (key type) values but are considered to be
// equivalent alternatives by the application using them.)  The "kid"
// value is a case-sensitive string.  Use of this member is OPTIONAL.
// When used with JWS or JWE, the "kid" value is used to match a JWS or
// JWE "kid" Header Parameter value." - RFC 7517
// https://datatracker.ietf.org/doc/html/rfc7517#section-4.5
func (f *PublicKeyFinder) ByKeyID(ctx context.Context, issuer string, keyID string) (*PublicKeyRecord, error) {
        jwks, err := f.fetchAndParseJwks(ctx, issuer)
        if err != nil {
                return nil, fmt.Errorf(`failed to fetch JWK set: %w`, err)
        }

        // If keyID is blank and there is only one key in the JWKS, return that key
        key, ok := jwks.LookupKeyID(keyID)
        if ok {
                return NewPublicKeyRecord(key, issuer)
        }

        return nil, fmt.Errorf("no matching public key found for kid %s", keyID)
}

func (f *PublicKeyFinder) ByJKT(ctx context.Context, issuer string, jkt string) (*PublicKeyRecord, error) {
        jwks, err := f.fetchAndParseJwks(ctx, issuer)
        if err != nil {
                return nil, err
        }

        it := jwks.Keys(ctx)
        for it.Next(ctx) {
                key := it.Pair().Value.(jwk.Key)
                jktOfKey, err := key.Thumbprint(crypto.SHA256)
                if err != nil {
                        return nil, fmt.Errorf("error computing Thumbprint of key in JWKS: %w", err)
                }
                jktOfKeyB64 := util.Base64EncodeForJWT(jktOfKey)
                if jkt == string(jktOfKeyB64) {
                        return NewPublicKeyRecord(key, issuer)
                }
        }

        return nil, fmt.Errorf("no matching public key found for jkt %s", jkt)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package discover

import (
        "context"
        "crypto"
        "encoding/json"

        "github.com/lestrrat-go/jwx/v2/jwk"
)

func MockGetJwksByIssuer(publicKeys []crypto.PublicKey, keyIDs []string, algs []string) (JwksFetchFunc, error) {
        // Create JWKS (JWK Set)
        jwks := jwk.NewSet()

        for i, publicKey := range publicKeys {
                jwkKey, err := jwk.PublicKeyOf(publicKey)
                if err != nil {
                        return nil, err
                }

                if err := jwkKey.Set(jwk.AlgorithmKey, algs[i]); err != nil {
                        return nil, err
                }

                if keyIDs != nil {
                        if err := jwkKey.Set(jwk.KeyIDKey, keyIDs[i]); err != nil {
                                return nil, err
                        }
                }

                // Put our jwk into a set
                if err := jwks.AddKey(jwkKey); err != nil {
                        return nil, err
                }
        }

        jwksJson, err := json.Marshal(jwks)
        if err != nil {
                return nil, err
        }

        return func(ctx context.Context, issuer string) ([]byte, error) {
                return jwksJson, nil
        }, nil
}

func MockGetJwksByIssuerOneKey(publicKey crypto.PublicKey, keyID string, alg string) (JwksFetchFunc, error) {
        // Create JWKS (JWK Set)
        jwkKey, err := jwk.PublicKeyOf(publicKey)
        if err != nil {
                return nil, err
        }

        if err := jwkKey.Set(jwk.AlgorithmKey, alg); err != nil {
                return nil, err
        }
        if err := jwkKey.Set(jwk.KeyIDKey, keyID); err != nil {
                return nil, err
        }

        // Put our jwk into a set
        jwks := jwk.NewSet()
        if err := jwks.AddKey(jwkKey); err != nil {
                return nil, err
        }

        jwksJson, err := json.Marshal(jwks)
        if err != nil {
                return nil, err
        }

        return func(ctx context.Context, issuer string) ([]byte, error) {
                return jwksJson, nil
        }, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gitlab_example

import (
        "context"
        "encoding/base64"
        "fmt"

        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/verifier"
)

type Opts struct {
        altOp providers.OpenIdProvider
}

func SignWithGitlab(opts ...Opts) error {
        var op providers.OpenIdProvider

        // If an alternative OP is provided, use that instead of the default.
        // Currently only used for testing where a mockOP is provided.
        if len(opts) > 0 && opts[0].altOp != nil {
                op = opts[0].altOp
        } else {
                // Creates OpenID Provider (OP) configuration, this will be used to request the ID Token from Gitlab
                op = providers.NewGitlabCiOpFromEnvironment("OPENPUBKEY_JWT")
        }

        // Creates a new OpenPubkey client
        opkClient, err := client.New(op)
        if err != nil {
                return err
        }

        // Generates a PK Token by authorizing to the OpenID Provider
        pkt, err := opkClient.Auth(context.Background())
        if err != nil {
                return err
        }

        // Serialize the PK Token to JSON so we can print it. Typically this
        // serialization of the PK Token would be sent with the signed message
        pktJson, err := pkt.MarshalJSON()
        if err != nil {
                return err
        }
        fmt.Println("pkt:", string(pktJson))

        pktCom, _ := pkt.Compact()
        b64pktCom := base64.StdEncoding.EncodeToString(pktCom)
        fmt.Println("pkt compact:", string(b64pktCom))

        // Create a verifier to check that the PK Token is well formed
        // The OPK client does this as well, but for the purposes of the
        // example we show how a relying party might verify a PK Token
        verifier, err := verifier.New(op)
        if err != nil {
                return err
        }

        // Verify the PK Token. We supply the OP (gitlab) we wish to verify against
        err = verifier.VerifyPKToken(context.Background(), pkt)
        if err != nil {
                return err
        }

        // Sign a message over the user's public key in the PK Token
        msg := []byte("All is discovered - flee at once")
        signedMsg, err := pkt.NewSignedMessage(msg, opkClient.GetSigner())
        if err != nil {
                return err
        }
        fmt.Println("signedMsg:", string(signedMsg))

        // Verify the signed message
        _, err = pkt.VerifySignedMessage(signedMsg)
        if err != nil {
                return err
        }

        fmt.Println("Success!")
        return nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package main

import (
        "context"
        "fmt"
        "os"

        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/cosigner"
        "github.com/openpubkey/openpubkey/examples/mfa/mfacosigner"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/verifier"
)

func main() {
        provider := providers.NewGoogleOp()

        cosignerProvider := client.CosignerProvider{
                Issuer:       "http://localhost:3003",
                CallbackPath: "/mfacallback",
        }

        if len(os.Args) < 2 {
                fmt.Printf("Example MFA Cosigner: command choices are: login, mfa")
                return
        }

        command := os.Args[1]
        switch command {
        case "login":

                opk, err := client.New(provider,
                        client.WithCosignerProvider(&cosignerProvider),
                )
                if err != nil {
                        fmt.Println(err)
                        return
                }

                pkt, err := opk.Auth(context.TODO())
                if err != nil {
                        fmt.Println(err)
                        return
                }
                fmt.Println("New PK token generated")

                // Verify our pktoken including the cosigner signature
                cosVerifier := cosigner.NewCosignerVerifier(cosignerProvider.Issuer, cosigner.CosignerVerifierOpts{})
                verifier, err := verifier.New(provider, verifier.WithCosignerVerifiers(cosVerifier))
                if err != nil {
                        fmt.Println(err)
                        return
                }
                if err := verifier.VerifyPKToken(context.TODO(), pkt); err != nil {
                        fmt.Println("Failed to verify PK token:", err)
                        os.Exit(1)
                } else {
                        fmt.Println("PK token verified successfully!")
                }

                os.Exit(0)
        case "mfa":
                rpID := "localhost"
                serverUri := "http://localhost:3003"
                rpOrigin := "http://localhost:3003"
                rpDisplayName := "OpenPubkey"
                _, err := mfacosigner.NewMfaCosignerHttpServer(serverUri, rpID, rpOrigin, rpDisplayName)
                if err != nil {
                        fmt.Println("error starting mfa server: ", err)
                        return
                }

        default:
                fmt.Println("Unrecognized command:", command)
                fmt.Printf("Example MFA Cosigner: command choices are: login, mfa")
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package jwks

import (
        "crypto"
        "crypto/ecdsa"
        "crypto/elliptic"
        "encoding/hex"
        "encoding/json"
        "fmt"
        "net"
        "net/http"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"
        "golang.org/x/crypto/sha3"
)

type JwksServer struct {
        uri       string
        jwksBytes []byte
}

// A very simple JWKS server for our MFA Cosigner example code.
func NewJwksServer(signer crypto.Signer, alg jwa.SignatureAlgorithm) (*JwksServer, string, error) {
        // Compute the kid (Key ID) as the SHA-3 of the public key
        pubkey := signer.Public().(*ecdsa.PublicKey) // TODO: handle non-ecdsa signers
        pubkeyBytes := elliptic.Marshal(pubkey, pubkey.X, pubkey.Y)
        pubkeyHash := sha3.Sum256(pubkeyBytes)
        kid := hex.EncodeToString(pubkeyHash[:])

        // Generate our JWKS using our signing key
        jwkKey, err := jwk.PublicKeyOf(signer)
        if err != nil {
                return nil, "", err
        }
        jwkKey.Set(jwk.AlgorithmKey, alg)
        jwkKey.Set(jwk.KeyIDKey, kid)

        // Put our jwk into a set
        keySet := jwk.NewSet()
        keySet.AddKey(jwkKey)

        // Now convert our key set into the raw bytes for printing later
        keySetBytes, _ := json.MarshalIndent(keySet, "", "  ")
        if err != nil {
                return nil, "", err
        }

        // Find an empty port
        listener, err := net.Listen("tcp", ":0")
        if err != nil {
                return nil, "", fmt.Errorf("failed to bind to an available port: %w", err)
        }

        server := &JwksServer{
                uri:       fmt.Sprintf("http://localhost:%d", listener.Addr().(*net.TCPAddr).Port),
                jwksBytes: keySetBytes,
        }

        // Host our JWKS at a localhost url
        mux := http.NewServeMux()
        mux.HandleFunc("/.well-known/jwks.json", server.printJWKS)
        go func() {
                http.Serve(listener, mux)
        }()
        return server, kid, nil
}

func (s *JwksServer) URI() string {
        return s.uri
}

func (s *JwksServer) printJWKS(w http.ResponseWriter, r *http.Request) {
        w.Header().Set("Content-Type", "application/json")
        w.Write(s.jwksBytes)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mfacosigner

import (
        "crypto"
        "crypto/rand"
        "fmt"

        "github.com/go-webauthn/webauthn/protocol"
        "github.com/go-webauthn/webauthn/webauthn"
        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/openpubkey/openpubkey/cosigner"
        "github.com/openpubkey/openpubkey/cosigner/mocks"
)

func NewUser(as *cosigner.AuthState) *user {
        return &user{
                id:          []byte(as.Sub),
                username:    as.Username,
                displayName: as.DisplayName,
        }
}

// This is intended as an example. Both sessionMap and users are not concurrency safe.
type MfaCosigner struct {
        *cosigner.AuthCosigner
        webAuthn   *webauthn.WebAuthn
        sessionMap map[string]*webauthn.SessionData
        users      map[cosigner.UserKey]*user
}

func New(signer crypto.Signer, alg jwa.SignatureAlgorithm, issuer, keyID string, cfg *webauthn.Config) (*MfaCosigner, error) {
        hmacKey := make([]byte, 64)
        if _, err := rand.Read(hmacKey); err != nil {
                return nil, err
        }

        wauth, err := webauthn.New(cfg)
        if err != nil {
                return nil, err
        }

        authCos, err := cosigner.New(signer, alg, issuer, keyID, mocks.NewAuthStateInMemoryStore(hmacKey))
        if err != nil {
                return nil, err
        }

        return &MfaCosigner{
                AuthCosigner: authCos,
                webAuthn:     wauth,
                sessionMap:   make(map[string]*webauthn.SessionData),
                users:        make(map[cosigner.UserKey]*user),
        }, nil
}

func (c *MfaCosigner) CheckIsRegistered(authID string) bool {
        authState, _ := c.AuthStateStore.LookupAuthState(authID)
        userKey := authState.UserKey()
        return c.IsRegistered(userKey)
}

func (c *MfaCosigner) IsRegistered(userKey cosigner.UserKey) bool {
        _, ok := c.users[userKey]
        return ok
}

func (c *MfaCosigner) BeginRegistration(authID string) (*protocol.CredentialCreation, error) {
        authState, _ := c.AuthStateStore.LookupAuthState(authID)
        userKey := authState.UserKey()

        if c.IsRegistered(userKey) {
                return nil, fmt.Errorf("already has a webauthn device registered for this user")
        }
        user := NewUser(authState)
        credCreation, session, err := c.webAuthn.BeginRegistration(user)
        if err != nil {
                return nil, err
        }
        c.sessionMap[authID] = session
        return credCreation, err
}

func (c *MfaCosigner) FinishRegistration(authID string, parsedResponse *protocol.ParsedCredentialCreationData) error {
        authState, _ := c.AuthStateStore.LookupAuthState(authID)
        session := c.sessionMap[authID]

        userKey := authState.UserKey()
        if c.IsRegistered(userKey) {
                return fmt.Errorf("already has a webauthn device registered for this user")
        }
        user := NewUser(authState)
        credential, err := c.webAuthn.CreateCredential(user, *session, parsedResponse)
        if err != nil {
                return err
        }
        user.AddCredential(*credential)

        // TODO: Should use some mechanism to ensure that a registration session
        // can't overwrite the result of another registration session for the same
        // user if the user interleaved their registration sessions. It is a very
        // unlikely possibility but it would be good to rule it out.
        c.users[userKey] = user
        return nil
}

func (c *MfaCosigner) BeginLogin(authID string) (*protocol.CredentialAssertion, error) {
        authState, _ := c.AuthStateStore.LookupAuthState(authID)
        userKey := authState.UserKey()

        if user, ok := c.users[userKey]; !ok {
                return nil, fmt.Errorf("user does not exist for userkey given %s", userKey)
        } else if credAssertion, session, err := c.webAuthn.BeginLogin(user); err != nil {
                return nil, err
        } else {
                c.sessionMap[authID] = session
                return credAssertion, err
        }

}

func (c *MfaCosigner) FinishLogin(authID string, parsedResponse *protocol.ParsedCredentialAssertionData) (string, string, error) {
        authState, _ := c.AuthStateStore.LookupAuthState(authID)
        session := c.sessionMap[authID]
        userKey := authState.UserKey()

        _, err := c.webAuthn.ValidateLogin(c.users[userKey], *session, parsedResponse)
        if err != nil {
                return "", "", err
        }

        if authcode, err := c.NewAuthcode(authID); err != nil {
                return "", "", err
        } else {
                return authcode, authState.RedirectURI, nil
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mocks

import (
        "crypto"
        "crypto/ecdsa"
        "crypto/rand"
        "crypto/sha256"

        "github.com/go-webauthn/webauthn/protocol"
        "github.com/go-webauthn/webauthn/protocol/webauthncbor"
        "github.com/go-webauthn/webauthn/protocol/webauthncose"
        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/openpubkey/openpubkey/util"
)

// For testing purposes we create a WebAuthn device to run the client part of the protocol
type WebAuthnDevice struct {
        signer     crypto.Signer
        PubkeyCbor []byte
        RpID       string
        RpIDHash   []byte
        Userhandle []byte
        RawID      []byte
        AuthFlags  byte
        Counter    uint32
}

func NewWebauthnDevice(rpID string) (*WebAuthnDevice, error) {
        alg := jwa.ES256
        signer, err := util.GenKeyPair(alg)
        if err != nil {
                return nil, err
        }

        pubkey := signer.Public().(*ecdsa.PublicKey)
        pubkeyCbor := webauthncose.EC2PublicKeyData{
                PublicKeyData: webauthncose.PublicKeyData{
                        KeyType:   int64(webauthncose.EllipticKey),
                        Algorithm: int64(webauthncose.AlgES256),
                },
                Curve:  int64(webauthncose.AlgES256),
                XCoord: pubkey.X.Bytes(),
                YCoord: pubkey.Y.Bytes(),
        }
        pubkeyCborBytes, err := webauthncbor.Marshal(pubkeyCbor)
        if err != nil {
                return nil, err
        }

        rpIDHash := sha256.Sum256([]byte(rpID))

        return &WebAuthnDevice{
                signer:     signer,
                PubkeyCbor: pubkeyCborBytes,
                RpID:       rpID,
                RpIDHash:   rpIDHash[:],

                // Checked by Webauthn RP to distinguish between different
                // users accounts sharing the same device with the same RP.
                // userHandle == user.WebAuthnID()?
                //
                // Not all devices can store a user handle it is allowed to be null
                // https://developer.mozilla.org/en-US/docs/Web/API/AuthenticatorAssertionResponse/userHandle
                //
                // In OpenPubkey MFA Cosigner RP we set this to the ID Token sub
                Userhandle: nil,

                // The ID of the public key credential held by the device
                RawID: []byte{5, 1, 1, 1, 1},

                //         Flag 0x41 has two bits set. 0b001 and 0b101
                //  FlagUserPresent Bit 00000001 - the user is present (UP flag)
                //  FlagUserVerified Bit 00000100 - user is verified using a biometric or PIN (UV flag)
                AuthFlags: 0x41,

                // Signature counter, used to identify cloned devices  see https://www.w3.org/TR/webauthn/#signature-counter
                Counter: 0,
        }, nil
}

func (wa *WebAuthnDevice) RegResp(createCreation *protocol.CredentialCreation) (*protocol.ParsedCredentialCreationData, error) {
        wa.Userhandle = []byte(createCreation.Response.User.ID.(protocol.URLEncodedBase64))

        return &protocol.ParsedCredentialCreationData{
                Response: protocol.ParsedAttestationResponse{
                        CollectedClientData: protocol.CollectedClientData{
                                Type:      protocol.CeremonyType("webauthn.create"),
                                Challenge: createCreation.Response.Challenge.String(),
                                Origin:    createCreation.Response.RelyingParty.ID,
                        },
                        AttestationObject: protocol.AttestationObject{
                                Format: "none",
                                AuthData: protocol.AuthenticatorData{
                                        RPIDHash: wa.RpIDHash,
                                        Counter:  wa.Counter,
                                        Flags:    protocol.AuthenticatorFlags(wa.AuthFlags),
                                        AttData: protocol.AttestedCredentialData{
                                                AAGUID:              make([]byte, 16),
                                                CredentialID:        wa.RawID,
                                                CredentialPublicKey: wa.PubkeyCbor,
                                        },
                                },
                        },
                        Transports: []protocol.AuthenticatorTransport{protocol.USB, protocol.NFC, "fake"},
                },
        }, nil
}

func (wa *WebAuthnDevice) LoginResp(credAssert *protocol.CredentialAssertion) (*protocol.ParsedCredentialAssertionData, error) {
        loginRespData := &protocol.ParsedCredentialAssertionData{
                ParsedPublicKeyCredential: protocol.ParsedPublicKeyCredential{
                        // Checked by Webauthn RP to see if public key supplied is on the
                        // allowlist of public keys for this user:
                        // parsedResponse.RawID == session.AllowedCredentialIDs?
                        RawID: wa.RawID,
                },
                Response: protocol.ParsedAssertionResponse{
                        CollectedClientData: protocol.CollectedClientData{
                                Type:      protocol.CeremonyType("webauthn.get"),
                                Challenge: credAssert.Response.Challenge.String(),
                                Origin:    wa.RpID,
                        },
                        AuthenticatorData: protocol.AuthenticatorData{
                                RPIDHash: wa.RpIDHash,
                                Counter:  wa.Counter,
                                Flags:    protocol.AuthenticatorFlags(wa.AuthFlags),
                        },
                        UserHandle: wa.Userhandle, // Not a required field:
                },
        }

        return wa.SignLoginChallenge(loginRespData)
}

func (wa *WebAuthnDevice) SignLoginChallenge(loginRespData *protocol.ParsedCredentialAssertionData) (*protocol.ParsedCredentialAssertionData, error) {
        clientDataHash := sha256.Sum256(loginRespData.Raw.AssertionResponse.ClientDataJSON)
        sigData := append(loginRespData.Raw.AssertionResponse.AuthenticatorData, clientDataHash[:]...)
        sigHash := sha256.Sum256(sigData)
        sigWebauthn, err := wa.signer.Sign(rand.Reader, sigHash[:], crypto.SHA256)
        if err != nil {
                return nil, err
        }
        loginRespData.Response.Signature = sigWebauthn
        return loginRespData, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mfacosigner

import (
        "encoding/json"
        "fmt"
        "net/http"

        "github.com/go-webauthn/webauthn/protocol"
        "github.com/go-webauthn/webauthn/webauthn"
        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/openpubkey/openpubkey/examples/mfa/mfacosigner/jwks"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/util"
)

type Server struct {
        cosigner *MfaCosigner
        jwksUri  string
}

func NewMfaCosignerHttpServer(serverUri, rpID, rpOrigin, RPDisplayName string) (*Server, error) {
        server := &Server{}

        // WebAuthn configuration
        cfg := &webauthn.Config{
                RPDisplayName: RPDisplayName,
                RPID:          rpID,
                RPOrigin:      rpOrigin,
        }

        // Generate the key pair for our cosigner
        alg := jwa.ES256
        signer, err := util.GenKeyPair(alg)
        if err != nil {
                return nil, err
        }

        jwksServer, kid, err := jwks.NewJwksServer(signer, alg)
        if err != nil {
                return nil, err
        }
        jwksHost := jwksServer.URI()
        server.jwksUri = fmt.Sprintf("%s/.well-known/jwks.json", jwksHost)
        issuer := rpOrigin

        fmt.Println("JWKS hosted at", server.jwksUri)
        server.cosigner, err = New(signer, alg, issuer, kid, cfg)
        if err != nil {
                return nil, err
        }

        mux := http.NewServeMux()
        mux.Handle("/", http.FileServer(http.Dir("mfacosigner/static")))
        mux.HandleFunc("/mfa-auth-init", server.initAuth)
        mux.HandleFunc("/check-registration", server.checkIfRegistered)
        mux.HandleFunc("/register/begin", server.beginRegistration)
        mux.HandleFunc("/register/finish", server.finishRegistration)
        mux.HandleFunc("/login/begin", server.beginLogin)
        mux.HandleFunc("/login/finish", server.finishLogin)
        mux.HandleFunc("/sign", server.signPkt)
        mux.HandleFunc("/.well-known/openid-configuration", server.wellKnownConf)

        err = http.ListenAndServe(":3003", mux)
        return server, err
}

func (s *Server) URI() string {
        return s.cosigner.Issuer
}

func (s *Server) initAuth(w http.ResponseWriter, r *http.Request) {
        err := r.ParseForm()
        if err != nil {
                return
        }
        pktB64 := []byte(r.URL.Query().Get("pkt"))
        pktJson, err := util.Base64DecodeForJWT(pktB64)
        if err != nil {
                return
        }
        var pkt *pktoken.PKToken
        if err := json.Unmarshal(pktJson, &pkt); err != nil {
                return
        }
        sig := []byte(r.URL.Query().Get("sig1"))

        authID, err := s.cosigner.InitAuth(pkt, sig)
        if err != nil {
                http.Error(w, "Error initiating authentication", http.StatusInternalServerError)
                return
        }
        mfapage := fmt.Sprintf("/?authid=%s", authID)

        http.Redirect(w, r, mfapage, http.StatusFound)
}

func (s *Server) checkIfRegistered(w http.ResponseWriter, r *http.Request) {
        authID, err := GetAuthID(r)
        if err != nil {
                http.Error(w, "Error in authID", http.StatusInternalServerError)
                return
        }
        registered := s.cosigner.CheckIsRegistered(authID)

        response, _ := json.Marshal(map[string]bool{
                "isRegistered": registered,
        })

        w.WriteHeader(200)
        w.Write(response)
}

func GetAuthID(r *http.Request) (string, error) {
        if err := r.ParseForm(); err != nil {
                return "", err
        }
        return string([]byte(r.URL.Query().Get("authid"))), nil
}

func (s *Server) beginRegistration(w http.ResponseWriter, r *http.Request) {
        authID, err := GetAuthID(r)
        if err != nil {
                http.Error(w, "Error in authID", http.StatusInternalServerError)
                return
        }

        options, err := s.cosigner.BeginRegistration(authID)

        optionsJson, err := json.Marshal(options)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        w.Header().Set("Content-Type", "application/json")
        w.Write(optionsJson)
}

func (s *Server) finishRegistration(w http.ResponseWriter, r *http.Request) {
        authID, err := GetAuthID(r)
        if err != nil {
                http.Error(w, "Error in authID", http.StatusInternalServerError)
                return
        }

        parsedResponse, err := protocol.ParseCredentialCreationResponse(r)
        if err != nil {
                http.Error(w, "Error in parsing credential", http.StatusInternalServerError)
                return
        }

        err = s.cosigner.FinishRegistration(authID, parsedResponse)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        w.WriteHeader(201)
        fmt.Println("MFA registration complete")
}

func (s *Server) beginLogin(w http.ResponseWriter, r *http.Request) {
        authID, err := GetAuthID(r)
        if err != nil {
                http.Error(w, "Error in authID", http.StatusInternalServerError)
                return
        }
        options, err := s.cosigner.BeginLogin(authID)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        optionsJson, err := json.Marshal(options)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        w.Header().Set("Content-Type", "application/json")
        w.Write(optionsJson)
}

func (s *Server) finishLogin(w http.ResponseWriter, r *http.Request) {
        authID, err := GetAuthID(r)
        if err != nil {
                http.Error(w, "Error in authID", http.StatusInternalServerError)
                return
        }

        parsedResponse, err := protocol.ParseCredentialRequestResponse(r)
        if err != nil {
                http.Error(w, "Error in parsing credential", http.StatusInternalServerError)
                return
        }

        authcode, ruri, err := s.cosigner.FinishLogin(authID, parsedResponse)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        redirectURIl := fmt.Sprintf("%s?authcode=%s", ruri, authcode)
        response, _ := json.Marshal(map[string]string{
                "redirect_uri": redirectURIl,
        })
        w.WriteHeader(201)
        w.Write(response)
}

func (s *Server) signPkt(w http.ResponseWriter, r *http.Request) {
        err := r.ParseForm()
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }

        sig := []byte(r.URL.Query().Get("sig2"))

        if cosSig, err := s.cosigner.RedeemAuthcode(sig); err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        } else {
                cosSigB64 := util.Base64EncodeForJWT(cosSig)
                w.WriteHeader(201)
                w.Write(cosSigB64)
        }
}

func (s *Server) wellKnownConf(w http.ResponseWriter, r *http.Request) {
        type WellKnown struct {
                Issuer  string `json:"issuer"`
                JwksUri string `json:"jwks_uri"`
        }

        wk := WellKnown{
                Issuer:  s.cosigner.Issuer,
                JwksUri: s.jwksUri,
        }

        wkJson, err := json.Marshal(wk)
        if err != nil {
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
        }
        w.WriteHeader(200)
        w.Write(wkJson)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mfacosigner

import (
        "github.com/go-webauthn/webauthn/webauthn"
)

type user struct {
        id          []byte
        username    string
        displayName string
        credentials []webauthn.Credential
}

var _ webauthn.User = (*user)(nil)

func (u *user) WebAuthnID() []byte {
        return u.id
}

func (u *user) WebAuthnName() string {
        return u.username
}

func (u *user) WebAuthnDisplayName() string {
        return u.displayName
}

func (u *user) WebAuthnIcon() string {
        return ""
}

func (u *user) AddCredential(cred webauthn.Credential) {
        u.credentials = append(u.credentials, cred)
}

func (u *user) WebAuthnCredentials() []webauthn.Credential {
        return u.credentials
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package main

import (
        "context"
        "fmt"

        "github.com/goccy/go-json"
        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/verifier"
)

func Sign(op client.OpenIdProvider) ([]byte, []byte, error) {
        // Create a OpenPubkey client, this automatically generates a fresh
        // key pair (public key, signing key). The public key is added to any
        // PK Tokens the client generates
        opkClient, err := client.New(op)
        if err != nil {
                return nil, nil, err
        }

        // Generate a PK Token by authenticating to the OP (Google)
        pkt, err := opkClient.Auth(context.Background())
        if err != nil {
                return nil, nil, err
        }

        // Use the signing key that the client just generated to sign the message
        msg := []byte("All is discovered - flee at once")
        signedMsg, err := pkt.NewSignedMessage(msg, opkClient.GetSigner())
        if err != nil {
                return nil, nil, err
        }

        // Serialize the PK Token as JSON and distribute it with the signed message
        pktJson, err := json.Marshal(pkt)
        if err != nil {
                return nil, nil, err
        }
        return pktJson, signedMsg, nil
}

func Verify(op client.OpenIdProvider, pktJson []byte, signedMsg []byte) error {
        // Create a PK Token object from the PK Token JSON
        pkt := new(pktoken.PKToken)
        err := json.Unmarshal(pktJson, &pkt)
        if err != nil {
                return err
        }

        // Verify that PK Token is issued by the OP you wish to use
        pktVerifier, err := verifier.New(op)
        if err != nil {
                return err
        }
        err = pktVerifier.VerifyPKToken(context.Background(), pkt)
        if err != nil {
                return err
        }

        // Check that the message verifies under the user's public key in the PK Token
        msg, err := pkt.VerifySignedMessage(signedMsg)
        if err != nil {
                return err
        }

        // Get the signer's email address from ID Token inside the PK Token
        idtClaims := new(oidc.OidcClaims)
        if err := json.Unmarshal(pkt.Payload, idtClaims); err != nil {
                return err
        }

        fmt.Printf("Verification successful: %s (%s) signed the message '%s'\n", idtClaims.Email, idtClaims.Issuer, string(msg))
        return nil
}

func main() {
        opOptions := providers.GetDefaultGoogleOpOptions()
        // Change this to true to turn on GQ signatures
        opOptions.GQSign = false
        op := providers.NewGoogleOpWithOptions(opOptions)
        pktJson, signedMsg, err := Sign(op)
        if err != nil {
                fmt.Println("Failed to sign message:", err)
                return
        }
        err = Verify(op, pktJson, signedMsg)
        if err != nil {
                fmt.Println("Failed to verify message:", err)
                return
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package main

import (
        "bytes"
        "context"
        "crypto"
        "crypto/ecdsa"
        "crypto/rand"
        "encoding/base64"
        "encoding/hex"
        "encoding/json"
        "fmt"
        "os"
        "os/signal"
        "path"
        "syscall"

        "github.com/awnumar/memguard"

        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/client/choosers"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/util"
        "github.com/openpubkey/openpubkey/verifier"
        "golang.org/x/crypto/sha3"
)

var (
        // File names for when we save or load our pktoken and the corresponding signing key
        skFileName  = "key.pem"
        pktFileName = "pktoken.json"
)

func main() {
        // Safely terminate in case of an interrupt signal
        memguard.CatchInterrupt()

        // Purge the session when we return
        defer memguard.Purge()

        if len(os.Args) < 2 {
                fmt.Printf("OpenPubkey: command choices are login, sign, and cert")
                return
        }

        gqSign := false

        // Directory for saving data
        outputDir := "output/google"

        command := os.Args[1]
        switch command {
        case "login":
                if err := login(outputDir, gqSign); err != nil {
                        fmt.Println("Error logging in:", err)
                } else {
                        fmt.Println("Login successful!")
                }
        case "sign":
                message := "sign me!!"
                if err := sign(message, outputDir); err != nil {
                        fmt.Println("Failed to sign test message:", err)
                }
        default:
                fmt.Println("Unrecognized command:", command)
        }
}

func login(outputDir string, gqSign bool) error {
        googleOpOptions := providers.GetDefaultGoogleOpOptions()
        googleOpOptions.GQSign = gqSign
        googleOp := providers.NewGoogleOpWithOptions(googleOpOptions)

        azureOpOptions := providers.GetDefaultAzureOpOptions()
        azureOpOptions.GQSign = gqSign
        azureOp := providers.NewAzureOpWithOptions(azureOpOptions)

        gitlabOpOptions := providers.GetDefaultGitlabOpOptions()
        gitlabOpOptions.GQSign = gqSign
        gitlabOp := providers.NewGitlabOpWithOptions(gitlabOpOptions)

        helloOpOptions := providers.GetDefaultHelloOpOptions()
        helloOpOptions.GQSign = gqSign
        helloOp := providers.NewHelloOpWithOptions(helloOpOptions)

        ctx, cancel := context.WithCancel(context.Background())
        sigs := make(chan os.Signal, 1)
        signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
        go func() {
                <-sigs
                fmt.Printf("Received shutdown signal, exiting... %v\n", sigs)
                cancel()
        }()

        openBrowser := true
        op, err := choosers.NewWebChooser(
                []providers.BrowserOpenIdProvider{googleOp, azureOp, helloOp, gitlabOp},
                openBrowser,
        ).ChooseOp(ctx)
        if err != nil {
                return err
        }

        opkClient, err := client.New(op)
        if err != nil {
                return err
        }

        pkt, err := opkClient.Auth(ctx,
                client.WithExtraClaim("extra", "yes"))
        if err != nil {
                return err
        }

        // Pretty print our json token
        pktJson, err := json.MarshalIndent(pkt, "", "  ")
        if err != nil {
                return err
        }
        fmt.Println(string(pktJson))
        pktCom, err := pkt.Compact()
        if err != nil {
                return err
        }
        fmt.Println("Compact", len(pktCom), string(pktCom))

        if opkClient.Op != helloOp {
                newPkt, err := opkClient.Refresh(ctx)
                if err != nil {
                        return err
                }
                fmt.Println("refreshed ID Token", string(newPkt.FreshIDToken))

                // Verify that PK Token is issued by the OP you wish to use and that it has a refreshed ID Token
                ops := []verifier.ProviderVerifier{googleOp, azureOp, gitlabOp}
                pktVerifier, err := verifier.NewFromMany(ops, verifier.RequireRefreshedIDToken())
                if err != nil {
                        return err
                }
                err = pktVerifier.VerifyPKToken(context.Background(), newPkt)
                if err != nil {
                        return err
                }

                // Save our signer and pktoken by writing them to a file
                return saveLogin(outputDir, opkClient.GetSigner().(*ecdsa.PrivateKey), newPkt)
        } else {
                // HelloOP does not support refresh tokens
                fmt.Println("skipping ID Token refresh for Hello OP as it does not support refresh tokens")

                ops := []verifier.ProviderVerifier{googleOp, azureOp, gitlabOp, helloOp}
                pktVerifier, err := verifier.NewFromMany(ops)
                if err != nil {
                        return err
                }
                err = pktVerifier.VerifyPKToken(context.Background(), pkt)
                if err != nil {
                        return err
                }

                // Save our signer and pktoken by writing them to a file
                return saveLogin(outputDir, opkClient.GetSigner().(*ecdsa.PrivateKey), pkt)
        }

}

func sign(message string, outputDir string) error {
        signer, pkt, err := loadLogin(outputDir)
        if err != nil {
                return fmt.Errorf("failed to load client state: %w", err)
        }

        msgHashSum := sha3.Sum256([]byte(message))
        sig, err := signer.Sign(rand.Reader, msgHashSum[:], crypto.SHA256)
        if err != nil {
                return err
        }

        fmt.Println("Signed Message:", message)
        fmt.Println("Praise Sigma:", base64.StdEncoding.EncodeToString(sig))
        fmt.Println("Hash:", hex.EncodeToString(msgHashSum[:]))
        fmt.Println("Cert:")

        pktJson, err := json.Marshal(pkt)
        if err != nil {
                return err
        }

        // Pretty print our json token
        var prettyJSON bytes.Buffer
        if err := json.Indent(&prettyJSON, pktJson, "", "  "); err != nil {
                return err
        }
        fmt.Println(prettyJSON.String())

        return nil
}

func saveLogin(outputDir string, sk *ecdsa.PrivateKey, pkt *pktoken.PKToken) error {
        if err := os.MkdirAll(outputDir, 0777); err != nil {
                return err
        }

        skFilePath := path.Join(outputDir, skFileName)
        if err := util.WriteSKFile(skFilePath, sk); err != nil {
                return err
        }

        pktFilePath := path.Join(outputDir, pktFileName)
        pktJson, err := json.Marshal(pkt)
        if err != nil {
                return err
        }
        return os.WriteFile(pktFilePath, pktJson, 0600)
}

func loadLogin(outputDir string) (crypto.Signer, *pktoken.PKToken, error) {
        skFilePath := path.Join(outputDir, skFileName)
        key, err := util.ReadSKFile(skFilePath)
        if err != nil {
                return nil, nil, err
        }

        pktFilePath := path.Join(outputDir, pktFileName)
        pktJson, err := os.ReadFile(pktFilePath)
        if err != nil {
                return nil, nil, err
        }

        var pkt *pktoken.PKToken
        if err := json.Unmarshal(pktJson, &pkt); err != nil {
                return nil, nil, err
        }

        return key, pkt, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package ca

import (
        "bytes"
        "context"
        "crypto/ecdsa"
        "crypto/elliptic"
        "crypto/rand"
        "crypto/x509"
        "crypto/x509/pkix"
        "encoding/json"
        "encoding/pem"
        "fmt"
        "math/big"
        "time"

        "github.com/lestrrat-go/jwx/v2/jwa"

        "github.com/openpubkey/openpubkey/cert"
        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/verifier"
)

type Ca struct {
        pksk *ecdsa.PrivateKey
        Alg  jwa.KeyAlgorithm
        // CaCertBytes []byte
        RootCertPem []byte
        op          client.OpenIdProvider
}

func New(op client.OpenIdProvider) (*Ca, error) {
        ca := Ca{
                op: op,
        }
        alg := string(jwa.ES256)
        err := ca.KeyGen(alg)
        if err != nil {
                return nil, err
        }

        return &ca, nil
}

func (a *Ca) KeyGen(alg string) error {
        a.Alg = jwa.KeyAlgorithmFrom(alg)

        pksk, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
        if err != nil {
                return err
        }
        a.pksk = pksk

        caTemplate := &x509.Certificate{
                SerialNumber: big.NewInt(2019),
                Subject: pkix.Name{
                        Organization:  []string{"Openpubkey-test-ca-cert"},
                        Country:       []string{"International"},
                        Province:      []string{""},
                        Locality:      []string{""},
                        StreetAddress: []string{"255 Test St."},
                        PostalCode:    []string{""},
                },
                NotBefore:             time.Now(),
                NotAfter:              time.Now().AddDate(10, 0, 0),
                IsCA:                  true,
                ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageCodeSigning},
                KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
                BasicConstraintsValid: true,
        }

        caBytes, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &a.pksk.PublicKey, a.pksk)
        if err != nil {
                return err
        }

        caPEM := new(bytes.Buffer)
        pem.Encode(caPEM, &pem.Block{
                Type:  "CERTIFICATE",
                Bytes: caBytes,
        })
        a.RootCertPem = caPEM.Bytes()

        return nil
}

func (a *Ca) CheckPKToken(pktJson []byte) (*pktoken.PKToken, error) {
        pkt := new(pktoken.PKToken)
        if err := json.Unmarshal(pktJson, pkt); err != nil {
                return nil, err
        }

        verifier, err := verifier.New(a.op)
        if err != nil {
                return nil, err
        }
        if err := verifier.VerifyPKToken(context.TODO(), pkt); err != nil {
                return nil, fmt.Errorf("failed to verify PK token: %w", err)
        }

        return pkt, nil
}

func (a *Ca) PktToSignedX509(pktJson []byte) ([]byte, error) {
        pkt, err := a.CheckPKToken(pktJson)
        if err != nil {
                return nil, err
        }

        pktUpk, err := ExtractRawPubkey(pkt)
        if err != nil {
                return nil, err
        }

        subTemplate, err := cert.PktToX509Template(pkt)
        if err != nil {
                return nil, err
        }

        rootCert, _ := pem.Decode(a.RootCertPem)
        if rootCert == nil {
                return nil, fmt.Errorf("failed to parse certificate PEM")
        }
        caTemplate, err := x509.ParseCertificate(rootCert.Bytes)
        if err != nil {
                return nil, err
        }

        subCertBytes, err := x509.CreateCertificate(rand.Reader, subTemplate, caTemplate, pktUpk, a.pksk)
        if err != nil {
                return nil, err
        }

        subCert, err := x509.ParseCertificate(subCertBytes)
        if err != nil {
                return nil, err
        }

        var pemSubCert bytes.Buffer
        err = pem.Encode(&pemSubCert, &pem.Block{Type: "CERTIFICATE", Bytes: subCert.Raw})
        if err != nil {
                return nil, err
        }
        return pemSubCert.Bytes(), nil
}

// VerifyPktCert checks that the X509 cert is signed by the CA and that
// the PK Token in the cert matches the public key in the cert.
func (a *Ca) VerifyPktCert(issuedCertPEM []byte) error {

        roots := x509.NewCertPool()
        ok := roots.AppendCertsFromPEM([]byte(a.RootCertPem))
        if !ok {
                return fmt.Errorf("failed to parse root certificate")
        }

        block, _ := pem.Decode([]byte(issuedCertPEM))
        if block == nil {
                return fmt.Errorf("failed to parse certificate PEM")
        }
        cert, err := x509.ParseCertificate(block.Bytes)
        if err != nil {
                return fmt.Errorf("failed to parse certificate PEM: %w", err)
        }
        _, err = cert.Verify(x509.VerifyOptions{
                Roots:         roots,
                Intermediates: x509.NewCertPool(),
                KeyUsages:     []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning},
        })
        if err != nil {
                return fmt.Errorf("failed to verify certificate: %w", err)
        }

        pktJson := cert.SubjectKeyId
        pkt := new(pktoken.PKToken)
        if err := json.Unmarshal(pktJson, pkt); err != nil {
                return err
        }
        pktUpk, err := ExtractRawPubkey(pkt)
        if err != nil {
                return err
        }

        certPublickey := cert.PublicKey.(*ecdsa.PublicKey)
        if !certPublickey.Equal(pktUpk) {
                return fmt.Errorf("public key in cert does not match PK Token's public key")
        }

        certPublicKeyBytes, err := x509.MarshalPKIXPublicKey(certPublickey)
        if err := json.Unmarshal(pktJson, pkt); err != nil {
                return err
        }
        if string(cert.RawSubjectPublicKeyInfo) != string(certPublicKeyBytes) {
                return fmt.Errorf("certificate raw subject public key info does not match ephemeral public key")
        }

        // Verification succeeds
        return nil
}

func ExtractRawPubkey(pkt *pktoken.PKToken) (interface{}, error) {
        cic, err := pkt.GetCicValues()
        if err != nil {
                return nil, err
        }
        upk := cic.PublicKey()
        var rawUpk interface{} // This is the raw key, like *rsa.PrivateKey or *ecdsa.PrivateKey
        if err := upk.Raw(&rawUpk); err != nil {
                return nil, err
        }
        return rawUpk, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package main

import (
        "context"
        "encoding/json"
        "fmt"
        "os"

        "github.com/awnumar/memguard"

        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/examples/x509/ca"
        "github.com/openpubkey/openpubkey/providers"
)

func main() {
        // Safely terminate in case of an interrupt signal
        memguard.CatchInterrupt()

        // Purge the session when we return
        defer memguard.Purge()

        if len(os.Args) < 2 {
                fmt.Printf("OpenPubkey: command choices are login, sign, and cert")
                return
        }

        command := os.Args[1]
        switch command {
        case "login":
                opOpts := providers.GetDefaultGoogleOpOptions()
                opOpts.GQSign = true
                op := providers.NewGoogleOp()
                if err := login(op); err != nil {
                        fmt.Println("Error logging in:", err)
                } else {
                        fmt.Println("Login and X509 issuance successful!")
                }
        default:
                fmt.Println("Unrecognized command:", command)
        }
}

func login(op client.OpenIdProvider) error {
        opkClient, err := client.New(
                op,
        )
        if err != nil {
                return err
        }

        pkt, err := opkClient.Auth(context.Background())
        if err != nil {
                return err
        }

        // Pretty print our json token
        pktJson, err := json.MarshalIndent(pkt, "", "  ")
        if err != nil {
                return err
        }
        CertAuth, err := ca.New(op)
        if err != nil {
                return err
        }

        pemSubCert, err := CertAuth.PktToSignedX509(pktJson)
        if err != nil {
                return err
        }
        fmt.Println("Issued Cert: \n", string(pemSubCert))

        msg := []byte("All is discovered - flee at once")
        signedMsg, err := pkt.NewSignedMessage(msg, opkClient.GetSigner())
        if err != nil {
                return err
        }
        println("Signed Message: \n", string(signedMsg))

        err = CertAuth.VerifyPktCert(pemSubCert)
        if err != nil {
                return err
        }

        return nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
        "math/big"

        "filippo.io/bigmod"
)

// leaks only the size of x
func modAsInt(x *bigmod.Modulus) *big.Int {
        return new(big.Int).SetBytes(x.Nat().Bytes(x))
}

// leaks only the size of x
func natAsInt(x *bigmod.Nat, m *bigmod.Modulus) *big.Int {
        return new(big.Int).SetBytes(x.Bytes(m))
}

// leaks only the size of x
func intAsNat(x *big.Int, m *bigmod.Modulus) (*bigmod.Nat, error) {
        return bigmod.NewNat().SetBytes(x.Bytes(), m)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
        "crypto/rsa"
        "fmt"
        "io"
        "math/big"

        "filippo.io/bigmod"
        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "golang.org/x/crypto/sha3"
)

var GQ256 = jwa.SignatureAlgorithm("GQ256")

func init() {
        jwa.RegisterSignatureAlgorithm(GQ256)
}

type OptsStruct struct {
        extraClaims map[string]any
}
type Opts func(a *OptsStruct)

// WithExtraClaim specifies additional values to be included in the
// GQ signed JWT. These claims will be included in the protected header
// of the JWT
// Example use:
//
//        WithExtraClaim("claimKey", "claimValue")
func WithExtraClaim(k string, v string) Opts {
        return func(a *OptsStruct) {
                if a.extraClaims == nil {
                        a.extraClaims = map[string]any{}
                }
                a.extraClaims[k] = v
        }
}

// GQ256SignJWT takes a rsaPublicKey and signed JWT and computes a GQ1 signature
// on the JWT. It returns a JWT whose RSA signature has been replaced by
// the GQ signature. It is wrapper around SignerVerifier.SignJWT
// an additional check that the correct rsa public key has been supplied.
// Use this instead of SignerVerifier.SignJWT.
func GQ256SignJWT(rsaPublicKey *rsa.PublicKey, jwt []byte, opts ...Opts) ([]byte, error) {
        _, err := jws.Verify(jwt, jws.WithKey(jwa.RS256, rsaPublicKey))
        if err != nil {
                return nil, fmt.Errorf("incorrect public key supplied when GQ signing jwt: %w", err)
        }
        sv, err := New256SignerVerifier(rsaPublicKey)
        if err != nil {
                return nil, fmt.Errorf("error creating GQ signer: %w", err)
        }
        gqJWT, err := sv.SignJWT(jwt, opts...)
        if err != nil {
                return nil, fmt.Errorf("error creating GQ signature: %w", err)
        }
        return gqJWT, nil
}

// GQ256VerifyJWT verifies a GQ1 signature over GQ signed JWT
func GQ256VerifyJWT(rsaPublicKey *rsa.PublicKey, gqToken []byte) (bool, error) {
        sv, err := New256SignerVerifier(rsaPublicKey)
        if err != nil {
                return false, fmt.Errorf("error creating GQ signer: %w", err)
        }
        return sv.VerifyJWT(gqToken), nil
}

// Signer allows for creating GQ1 signatures messages.
type Signer interface {
        // Sign creates a GQ1 signature over the given message with the given GQ1 private number.
        Sign(private []byte, message []byte) ([]byte, error)
        // SignJWT creates a GQ1 signature over the JWT token's header/payload with a GQ1 private number derived from the JWT signature.
        //
        // This works because a GQ1 private number can be calculated as the inverse mod n of an RSA signature, where n is the public RSA modulus.
        SignJWT(jwt []byte, opts ...Opts) ([]byte, error)
}

// Signer allows for verifying GQ1 signatures.
type Verifier interface {
        // Verify verifies a GQ1 signature over a message, using the public identity of the signer.
        Verify(signature []byte, identity []byte, message []byte) bool

        // Compatible with SignJWT, this function verifies the GQ1 signature of the presented JSON Web Token.
        VerifyJWT(jwt []byte) bool
}

// SignerVerifier combines the Signer and Verifier interfaces.
type SignerVerifier interface {
        Signer
        Verifier
}

type signerVerifier struct {
        // n is the RSA public modulus (what Go's RSA lib calls N)
        n *bigmod.Modulus
        // v is the RSA public exponent (what Go's RSA lib calls E)
        v *big.Int
        // nBytes is the length of n in bytes
        nBytes int
        // vBytes is the length of v in bytes
        vBytes int
        // t is the signature length parameter
        t int
}

// Creates a new SignerVerifier specifically for GQ256, meaning the security parameter is 256.
func New256SignerVerifier(publicKey *rsa.PublicKey) (SignerVerifier, error) {
        return NewSignerVerifier(publicKey, 256)
}

// NewSignerVerifier creates a SignerVerifier from the RSA public key of the trusted third-party which creates
// the GQ1 private numbers.
//
// The securityParameter parameter is the level of desired security in bits. 256 is recommended.
func NewSignerVerifier(publicKey *rsa.PublicKey, securityParameter int) (SignerVerifier, error) {
        if publicKey.E != 65537 {
                // Danger: Currently it is unsafe to use this library with a RSA exponent other than 65537.
                // This issue is being tracked in https://github.com/openpubkey/openpubkey/issues/230
                return nil, fmt.Errorf("only 65537 is currently supported, unsupported RSA public key exponent: %d", publicKey.E)
        }
        n, v, nBytes, vBytes, err := parsePublicKey(publicKey)
        t := securityParameter / (vBytes * 8)

        return &signerVerifier{n, v, nBytes, vBytes, t}, err
}

func parsePublicKey(publicKey *rsa.PublicKey) (n *bigmod.Modulus, v *big.Int, nBytes int, vBytes int, err error) {
        n, err = bigmod.NewModulusFromBig(publicKey.N)
        if err != nil {
                return
        }
        v = big.NewInt(int64(publicKey.E))
        nLen := n.BitLen()
        vLen := v.BitLen() - 1 // note the -1; GQ1 only ever uses the (length of v) - 1, so we can just do this here rather than throughout
        nBytes = bytesForBits(nLen)
        vBytes = bytesForBits(vLen)
        return
}

func bytesForBits(bits int) int {
        return (bits + 7) / 8
}

var hash = func(byteCount int, data ...[]byte) ([]byte, error) {
        rng := sha3.NewShake256()
        for _, d := range data {
                rng.Write(d)
        }

        return randomBytes(rng, byteCount)
}

func randomBytes(rng io.Reader, byteCount int) ([]byte, error) {
        bytes := make([]byte, byteCount)

        _, err := io.ReadFull(rng, bytes)
        if err != nil {
                return nil, err
        }

        return bytes, nil
}

func OriginalJWTHeaders(jwt []byte) ([]byte, error) {
        token, err := jws.Parse(jwt)
        if err != nil {
                return nil, err
        }

        // a JWT is guaranteed to have exactly one signature
        headers := token.Signatures()[0].ProtectedHeaders()

        if headers.Algorithm() != GQ256 {
                return nil, fmt.Errorf("expected GQ256 alg, got %s", headers.Algorithm())
        }

        origHeaders := []byte(headers.KeyID())
        return origHeaders, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
        "crypto"
        "crypto/sha256"
)

// Hardcoded padding prefix for SHA-256 from https://github.com/golang/go/blob/eca5a97340e6b475268a522012f30e8e25bb8b8f/src/crypto/rsa/pkcs1v15.go#L268
var prefix = []byte{0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20}

// encodePKCS1v15 is taken from the go stdlib, see [crypto/rsa.SignPKCS1v15].
//
// https://github.com/golang/go/blob/eca5a97340e6b475268a522012f30e8e25bb8b8f/src/crypto/rsa/pkcs1v15.go#L287-L317
var encodePKCS1v15 = func(k int, data []byte) []byte {
        hashLen := crypto.SHA256.Size()
        tLen := len(prefix) + hashLen

        // EM = 0x00 || 0x01 || PS || 0x00 || T
        em := make([]byte, k)
        em[1] = 1
        for i := 2; i < k-tLen-1; i++ {
                em[i] = 0xff
        }
        copy(em[k-tLen:k-hashLen], prefix)

        hashed := sha256.Sum256(data)
        copy(em[k-hashLen:k], hashed[:])
        return em
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
        "crypto/rand"
        "encoding/json"
        "fmt"
        "math/big"

        "filippo.io/bigmod"
        "github.com/awnumar/memguard"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/util"
)

// Sign creates a GQ1 signature over the given message with the given GQ1 private number.
//
// Comments throughout refer to stages as specified in the ISO/IEC 14888-2 standard.
func (sv *signerVerifier) Sign(private []byte, message []byte) ([]byte, error) {
        n, v, t := sv.n, sv.v, sv.t
        vBytes := sv.vBytes

        M := message
        Q, err := bigmod.NewNat().SetBytes(private, n)
        if err != nil {
                return nil, err
        }

        // Stage 1 - select t numbers, each consisting of nBytes random bytes.
        // In order to guarantee our operation is constant time, we deviate slightly
        // from the standard and directly select an integer less than n
        r, err := randomNumbers(t, sv.n)
        if err != nil {
                return nil, err
        }

        // Stage 2 - calculate test number W
        // for i from 1 to t, compute W_i <- r_i^v mod n
        // combine to form W
        var W []byte
        for i := 0; i < t; i++ {
                W_i := bigmod.NewNat().Exp(r[i], v.Bytes(), n)
                W = append(W, W_i.Bytes(n)...)
        }

        // Stage 3 - calculate question number R
        // hash W and M and take first t*vBytes bytes as R
        R, err := hash(t*vBytes, W, M)
        if err != nil {
                return nil, err
        }

        // split R into t numbers each consisting of vBytes bytes
        Rs := make([]*bigmod.Nat, t)
        for i := 0; i < t; i++ {
                Rs[i], err = new(bigmod.Nat).SetBytes(R[i*vBytes:(i+1)*vBytes], n)
                if err != nil {
                        return nil, err
                }
        }

        // Stage 4 - calculate witness number S
        // for i from 1 to t, compute S_i <- r_i * Q^{R_i} mod n
        // combine to form S
        var S []byte
        for i := 0; i < t; i++ {
                S_i := bigmod.NewNat().Exp(Q, Rs[i].Bytes(n), n)
                S_i.Mul(r[i], n)
                S = append(S, S_i.Bytes(n)...)
        }

        // proof is combination of R and S
        return encodeProof(R, S), nil
}

func (sv *signerVerifier) SignJWT(jwt []byte, opts ...Opts) ([]byte, error) {
        options := &OptsStruct{}
        for _, applyOpt := range opts {
                applyOpt(options)
        }
        // Ensure that someone doesn't use a reserved protected header claim name
        for _, reserved := range []string{"alg", "typ", "kid"} {
                if _, ok := options.extraClaims[reserved]; ok {
                        return nil, fmt.Errorf("use of reserved header name, %s, in additional headers", reserved)
                }
        }

        origHeaders, payload, signature, err := jws.SplitCompact(jwt)
        if err != nil {
                return nil, err
        }

        signingPayload := util.JoinJWTSegments(origHeaders, payload)

        headers := jws.NewHeaders()
        err = headers.Set(jws.AlgorithmKey, GQ256)
        if err != nil {
                return nil, err
        }
        err = headers.Set(jws.TypeKey, "JWT")
        if err != nil {
                return nil, err
        }
        err = headers.Set(jws.KeyIDKey, string(origHeaders))
        if err != nil {
                return nil, err
        }

        for k, v := range options.extraClaims {
                if err = headers.Set(k, v); err != nil {
                        return nil, err
                }
        }

        headersJSON, err := json.Marshal(headers)
        if err != nil {
                return nil, err
        }

        headersEnc := util.Base64EncodeForJWT(headersJSON)

        // When jwt is parsed it's split into base64-encoded bytes, but
        // we need the raw signature to calculate mod inverse
        decodedSig, err := util.Base64DecodeForJWT(signature)
        if err != nil {
                return nil, err
        }

        // GQ1 private number (Q) is inverse of RSA signature mod n
        private, err := sv.modInverse(memguard.NewBufferFromBytes(decodedSig))
        if err != nil {
                return nil, err
        }

        defer private.Destroy()

        gqSig, err := sv.Sign(private.Bytes(), signingPayload)
        if err != nil {
                return nil, err
        }

        // Now make a new GQ-signed token
        gqToken := util.JoinJWTSegments(headersEnc, payload, gqSig)

        return gqToken, nil
}

// modInverse finds the modular multiplicative inverse of the value stored in b
//
// All operations involving the secret value are performed either with constant-
// time methods or with blinding (if sv has a source of randomness)
func (sv *signerVerifier) modInverse(b *memguard.LockedBuffer) (*memguard.LockedBuffer, error) {
        x, err := bigmod.NewNat().SetBytes(b.Bytes(), sv.n)
        if err != nil {
                return nil, err
        }

        nInt := natAsInt(sv.n.Nat(), sv.n)
        var r *big.Int
        var rConstant, xr *bigmod.Nat

        // Apply RSA blinding to the ModInverse operation.
        // Translates the technique formerly used in the Go Standard Library before they
        // switched to bigmod in late 2022. Since bigmod does not yet support constant-time
        // ModInverse, we perform the blinding so that the value of the private key is not
        // detectable via side channel.
        // Ref: https://github.com/golang/go/blob/5f60f844beb0581a19cb425a3338d79d322a7db2/src/crypto/rsa/rsa.go#L567-L596
        //
        // For a secret value x, the idea is to find m = 1/x mod n by calculating
        // rm/r mod n ==> r/(xr) mod n, where r is a random value

        for {
                // draw r
                r, err = rand.Int(rand.Reader, nInt)
                if err != nil {
                        return nil, err
                }

                // compute xr = x * r
                xr, err = intAsNat(r, sv.n)
                if err != nil {
                        return nil, err
                }
                xr.Mul(x, sv.n)

                // check that xr has a multiplicative inverse mod n. It is exceedingly
                // rare but technically possible for it not to, in which case we need
                // to draw a new value for r
                xrInt := natAsInt(xr, sv.n)
                inverse := new(big.Int).ModInverse(xrInt, nInt)
                if inverse != nil {
                        break
                }
        }

        // overwrite x with the blinded value
        x = xr

        // calculate m/r mod n
        m := natAsInt(x, sv.n).ModInverse(natAsInt(x, sv.n), nInt)
        mConstant, err := intAsNat(m, sv.n)
        if err != nil {
                return nil, err
        }

        // remove the blinding by multiplying m/r by r
        rConstant, err = intAsNat(r, sv.n)
        if err != nil {
                return nil, err
        }
        mConstant.Mul(rConstant, sv.n)

        mFinal := natAsInt(mConstant, sv.n)

        // need to allocate memory for fixed length slice using FillBytes
        ret := make([]byte, len(b.Bytes()))
        defer b.Destroy()

        return memguard.NewBufferFromBytes(mFinal.FillBytes(ret)), nil
}

func encodeProof(R, S []byte) []byte {
        var bin []byte

        bin = append(bin, R...)
        bin = append(bin, S...)

        return util.Base64EncodeForJWT(bin)
}

var randomNumbers = func(t int, n *bigmod.Modulus) ([]*bigmod.Nat, error) {
        nInt := modAsInt(n)
        ys := make([]*bigmod.Nat, t)

        for i := 0; i < t; i++ {
                r, err := rand.Int(rand.Reader, nInt)
                if err != nil {
                        return nil, err
                }

                ys[i], err = intAsNat(r, n)
                if err != nil {
                        return nil, err
                }
        }

        return ys, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package gq

import (
        "bytes"
        "fmt"
        "math/big"

        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/util"
)

// Verify verifies a GQ1 signature over a message, using the public identity of the signer.
//
// Comments throughout refer to stages as specified in the ISO/IEC 14888-2 standard.
func (sv *signerVerifier) Verify(proof []byte, identity []byte, message []byte) bool {
        n, v, t := modAsInt(sv.n), sv.v, sv.t
        nBytes, vBytes := sv.nBytes, sv.vBytes

        M := message

        // Stage 0 - reject proof if it's the wrong size based on t
        R, S, err := sv.decodeProof(proof)
        if err != nil {
                return false
        }

        // Stage 1 - create public number G
        // currently this hardcoded to use PKCS#1 v1.5 padding as the format mechanism
        paddedIdentity := encodePKCS1v15(nBytes, identity)
        G := new(big.Int).SetBytes(paddedIdentity)

        // Stage 2 - parse signature numbers and recalculate test number W*
        // split R into t strings, each consisting of vBytes bytes
        Rs := make([]*big.Int, t)
        for i := 0; i < t; i++ {
                Rs[i] = new(big.Int).SetBytes(R[i*vBytes : (i+1)*vBytes])
        }

        // split S into t strings, each consisting of nBytes bytes
        Ss := make([]*big.Int, t)
        for i := 0; i < t; i++ {
                s_i := new(big.Int).SetBytes(S[i*nBytes : (i+1)*nBytes])
                // reject if S_i = 0 or >= n
                if s_i.Cmp(big.NewInt(0)) == 0 || s_i.Cmp(n) != -1 {
                        return false
                }
                Ss[i] = s_i
        }

        // recalculate test number W*
        // for i from 1 to t, compute W*_i <- S_i^v * G^{R_i} mod n
        // combine to form W*

        var Wstar []byte
        for i := 0; i < t; i++ {
                l := new(big.Int).Exp(Ss[i], v, n)
                r := new(big.Int).Exp(G, Rs[i], n)
                Wstar_i := new(big.Int).Mul(l, r)
                Wstar_i.Mod(Wstar_i, n)
                b := make([]byte, nBytes)
                Wstar = append(Wstar, Wstar_i.FillBytes(b)...)
        }

        // Stage 3 - recalculate question number R*
        // hash W* and M and take first t*vBytes bytes as R*
        Rstar, err := hash(t*vBytes, Wstar, M)
        if err != nil {
                // TODO: this can only happen if there's some error reading /dev/urandom or something
                // so should we return the proper error?
                return false
        }

        // Stage 4 - accept or reject depending on whether R and R* are identical
        return bytes.Equal(R, Rstar)
}

func (sv *signerVerifier) VerifyJWT(jwt []byte) bool {
        origHeaders, err := OriginalJWTHeaders(jwt)
        if err != nil {
                return false
        }

        _, payload, signature, err := jws.SplitCompact(jwt)
        if err != nil {
                return false
        }

        signingPayload := util.JoinJWTSegments(origHeaders, payload)

        return sv.Verify(signature, signingPayload, signingPayload)
}

func (sv *signerVerifier) decodeProof(s []byte) (R, S []byte, err error) {
        bin, err := util.Base64DecodeForJWT(s)
        if err != nil {
                return nil, nil, err
        }

        rSize := sv.vBytes * sv.t
        sSize := sv.nBytes * sv.t

        if len(bin) != rSize+sSize {
                return nil, nil, fmt.Errorf("not the correct size")
        }

        R = bin[:rSize]
        S = bin[rSize:]

        return R, S, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package oidc

import (
        "fmt"
)

type Jws struct {
        Payload    string      `json:"payload"`    // Base64 encoded
        Signatures []Signature `json:"signatures"` // Base64 encoded
}

type SigOptStruct struct {
        PublicHeader map[string]any
}
type SigOpts func(a *SigOptStruct)

// WithPublicHeader species that a public header be included in the
// signature. Public headers aren't Base64 encoded because they aren't signed.
// Example use: WithPublicHeader(map[string]any{"key1": "abc", "key2": "def"})
func WithPublicHeader(publicHeader map[string]any) SigOpts {
        return func(o *SigOptStruct) {
                o.PublicHeader = publicHeader
        }
}

func (j *Jws) AddSignature(token []byte, opts ...SigOpts) error {
        sigOpts := &SigOptStruct{}
        for _, applyOpt := range opts {
                applyOpt(sigOpts)
        }

        protected, payload, signature, err := SplitCompact(token)
        if err != nil {
                return err
        }
        if j.Payload != string(payload) {
                return fmt.Errorf("payload in compact token does not match existing payload in jws, expected=(%s), got=(%s)",
                        string(j.Payload),
                        string(payload))
        }
        sig := Signature{
                Protected: string(protected),
                Public:    sigOpts.PublicHeader,
                Signature: string(signature),
        }

        if j.Signatures == nil {
                j.Signatures = []Signature{}
        }
        j.Signatures = append(j.Signatures, sig)

        return nil
}

func (j *Jws) GetToken(i int) ([]byte, error) {
        if i < len(j.Signatures) && i >= 0 {
                return []byte(j.Signatures[i].Protected + "." + j.Payload + "." + j.Signatures[i].Signature), nil
        } else {
                return nil, fmt.Errorf("no signature at index i (%d), len(signatures) (%d)", i, len(j.Signatures))
        }
}

func (j *Jws) GetTokenByTyp(typ string) ([]byte, error) {
        matchingTokens := []Signature{}
        for _, v := range j.Signatures {
                if typFound, err := v.GetTyp(); err != nil {
                        return nil, err
                } else {
                        // Both the JWS standard and the OIDC standard states that typ is case sensitive
                        // so we treat it as case sensitive as well
                        //
                        // "The typ (type) header parameter is used to declare the type of the
                        // signed content. The typ value is case sensitive."
                        // https://openid.net/specs/draft-jones-json-web-signature-04.html#ReservedHeaderParameterName
                        //
                        // "The "typ" (type) Header Parameter is used by JWS applications to
                        // declare the media type [IANA.MediaTypes] of this complete JWS.
                        // [..] Per RFC 2045 [RFC2045], all media type values, subtype values, and
                        // parameter names are case insensitive. However, parameter values are case
                        // sensitive unless otherwise specified for the specific parameter."
                        // https://datatracker.ietf.org/doc/html/rfc7515#section-4.1.9
                        if typFound == typ {
                                matchingTokens = append(matchingTokens, v)
                        }
                }
        }
        if len(matchingTokens) > 1 {
                // Currently we only have one token per token typ. We can change this later
                // for COS tokens. This check prevents hidden tokens, where one token of
                // the same typ hides another token of the same typ.
                return nil, fmt.Errorf("more than one token found, all current token typs are unique")
        } else if len(matchingTokens) == 0 {
                // if typ not found return nil
                return nil, nil
        } else {
                return []byte(matchingTokens[0].Protected + "." + j.Payload + "." + matchingTokens[0].Signature), nil
        }
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package oidc

import (
        "fmt"
)

type Jwt struct {
        payload       string
        payloadClaims *OidcClaims
        signature     *Signature
        raw           []byte
}

func NewJwt(token []byte) (*Jwt, error) {
        protected, payload, signature, err := SplitCompact(token)
        if err != nil {
                return nil, err
        }
        idt := &Jwt{
                payload: string(payload),
                signature: &Signature{
                        Protected: string(protected),
                        Signature: string(signature),
                },
                raw: token,
        }
        if err := ParseJWTSegment(protected, &idt.signature.protectedClaims); err != nil {
                return nil, fmt.Errorf("error parsing protected: %w", err)
        }
        if err := ParseJWTSegment(payload, &idt.payloadClaims); err != nil {
                return nil, fmt.Errorf("error parsing payload: %w", err)
        }
        return idt, nil
}

func (i *Jwt) GetClaims() *OidcClaims {
        return i.payloadClaims
}

func (i *Jwt) GetPayload() string {
        return i.payload
}

func (i *Jwt) GetSignature() *Signature {
        return i.signature
}

func (i *Jwt) GetRaw() []byte {
        return i.raw
}

// Compares two JWTs and determines if they are for the same identity (subject)
func SameIdentity(t1, t2 []byte) error {
        token1, err := NewJwt(t1)
        if err != nil {
                return err
        }
        token2, err := NewJwt(t2)
        if err != nil {
                return err
        }

        // Subject identity can only be established within the same issuer
        if token1.GetClaims().Issuer != token2.GetClaims().Issuer {
                return fmt.Errorf("tokens have different issuers")
        }

        if token1.GetClaims().Subject != token2.GetClaims().Subject {
                return fmt.Errorf("token have a different subject claims")
        }
        return nil
}

// RequireOlder returns an error if t1 is not older than t2
func RequireOlder(t1, t2 []byte) error {
        token1, err := NewJwt(t1)
        if err != nil {
                return err
        }
        token2, err := NewJwt(t2)
        if err != nil {
                return err
        }

        // Check which token was issued first
        if token1.GetClaims().IssuedAt > token2.GetClaims().IssuedAt {
                return fmt.Errorf("tokens not issued in correct order")
        }
        return nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package oidc

import (
        "bytes"
        "encoding/json"
        "fmt"
        "strings"

        "github.com/openpubkey/openpubkey/util"
)

type OidcClaims struct {
        Issuer     string `json:"iss"`
        Subject    string `json:"sub"`
        Audience   string `json:"-"`
        Expiration int64  `json:"exp"`
        IssuedAt   int64  `json:"iat"`
        Email      string `json:"email,omitempty"`
        Nonce      string `json:"nonce,omitempty"`
        Username   string `json:"preferred_username,omitempty"`
        FirstName  string `json:"given_name,omitempty"`
        LastName   string `json:"family_name,omitempty"`
}

// Implement UnmarshalJSON for custom handling during JSON unmarshalling
func (id *OidcClaims) UnmarshalJSON(data []byte) error {
        // unmarshal audience claim separately to account for []string, https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.3
        type Alias OidcClaims
        aux := &struct {
                Audience any `json:"aud"`
                *Alias
        }{
                Alias: (*Alias)(id),
        }

        if err := json.Unmarshal(data, &aux); err != nil {
                return err
        }

        switch t := aux.Audience.(type) {
        case string:
                id.Audience = t
        case []any:
                audList := []string{}
                for _, v := range t {
                        audList = append(audList, v.(string))
                }
                id.Audience = strings.Join(audList, ",")
        default:
                id.Audience = ""
        }

        return nil
}

// SplitCompact splits a JWT and returns its three parts
// separately: protected headers, payload and signature.
// This is copied from github.com/lestrrat-go/jwx/v2/jws.SplitCompact
// We include it here so so that jwx is not a dependency of simpleJws
func SplitCompact(src []byte) ([]byte, []byte, []byte, error) {
        parts := bytes.Split(src, []byte("."))
        if len(parts) != 3 {
                return nil, nil, nil, fmt.Errorf(`invalid number of segments`)
        }
        return parts[0], parts[1], parts[2], nil
}

func ParseJWTSegment(segment []byte, v any) error {
        segmentJSON, err := util.Base64DecodeForJWT(segment)
        if err != nil {
                return fmt.Errorf("error decoding segment: %w", err)
        }

        err = json.Unmarshal(segmentJSON, v)
        if err != nil {
                return fmt.Errorf("error parsing segment: %w", err)
        }

        return nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package oidc

import (
        "encoding/json"

        "github.com/openpubkey/openpubkey/util"
)

type Signature struct {
        Protected       string                 `json:"protected"` // Base64 encoded
        protectedClaims *ProtectedClaims       // Unmarshalled protected claims
        Public          map[string]interface{} `json:"header,omitempty"`
        Signature       string                 `json:"signature"` // Base64 encoded
}

type ProtectedClaims struct {
        Alg   string `json:"alg"`
        Jkt   string `json:"jkt,omitempty"`
        KeyID string `json:"kid,omitempty"`
        Type  string `json:"typ,omitempty"`
        CIC   string `json:"cic,omitempty"`
}

func (s *Signature) GetTyp() (string, error) {
        decodedProtected, err := util.Base64DecodeForJWT([]byte(s.Protected))
        if err != nil {
                return "", err
        }
        type protectedTyp struct {
                Typ string `json:"typ"`
        }
        var ph protectedTyp
        err = json.Unmarshal(decodedProtected, &ph)
        if err != nil {
                return "", err
        }
        return ph.Typ, nil
}

func (s *Signature) GetProtectedClaims() *ProtectedClaims {
        return s.protectedClaims
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package clientinstance

import (
        "crypto"
        "crypto/rand"
        "encoding/hex"
        "encoding/json"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/util"
)

// Client Instance Claims, referred also as "cic" in the OpenPubKey paper
type Claims struct {
        publicKey jwk.Key

        // Claims are stored in the protected header portion of JWS signature
        protected map[string]any
}

// Client instance claims must relate to a single key pair
func NewClaims(publicKey jwk.Key, claims map[string]any) (*Claims, error) {
        // Make sure our JWK has the algorithm header set
        if publicKey.Algorithm().String() == "" {
                return nil, fmt.Errorf("user JWK requires algorithm to be set")
        }

        // Make sure no claims are using our reserved values
        for _, reserved := range []string{"alg", "upk", "rz", "typ"} {
                if _, ok := claims[reserved]; ok {
                        return nil, fmt.Errorf("use of reserved header name, %s, in additional headers", reserved)
                }
        }

        rand, err := generateRand()
        if err != nil {
                return nil, fmt.Errorf("failed to generate random value: %w", err)
        }

        // Assign required values
        claims["typ"] = "CIC"
        claims["alg"] = publicKey.Algorithm().String()
        claims["upk"] = publicKey
        claims["rz"] = rand

        return &Claims{
                publicKey: publicKey,
                protected: claims,
        }, nil
}

func ParseClaims(protected map[string]any) (*Claims, error) {
        // Get our standard headers and make sure they match up
        if _, ok := protected["rz"]; !ok {
                return nil, fmt.Errorf(`missing required "rz" claim`)
        }
        upk, ok := protected["upk"]
        if !ok {
                return nil, fmt.Errorf(`missing required "upk" claim`)
        }
        upkBytes, err := json.Marshal(upk)
        if err != nil {
                return nil, err
        }
        upkjwk, err := jwk.ParseKey(upkBytes)
        if err != nil {
                return nil, err
        }
        alg, ok := protected["alg"]
        if !ok {
                return nil, fmt.Errorf(`missing required "alg" claim`)
        } else if alg != upkjwk.Algorithm() {
                return nil, fmt.Errorf(`provided "alg" value different from algorithm provided in "upk" jwk`)
        }
        return &Claims{
                publicKey: upkjwk,
                protected: protected,
        }, nil
}

func (c *Claims) PublicKey() jwk.Key {
        return c.publicKey
}

func (c *Claims) KeyAlgorithm() jwa.KeyAlgorithm {
        return c.publicKey.Algorithm()
}

// Returns a hash of all client instance claims which includes a random value
func (c *Claims) Hash() ([]byte, error) {
        buf, err := json.Marshal(c.protected)
        if err != nil {
                return nil, err
        }

        return util.B64SHA3_256(buf), nil
}

// This function signs the payload of the provided token with the protected headers
// as defined by the client instance claims and returns a jwt in compact form.
func (c *Claims) Sign(signer crypto.Signer, algorithm jwa.KeyAlgorithm, token []byte) ([]byte, error) {
        _, payload, _, err := jws.SplitCompact(token)
        if err != nil {
                return nil, err
        }

        // We need to make sure we're signing the decoded bytes
        payloadDecoded, err := util.Base64DecodeForJWT(payload)
        if err != nil {
                return nil, err
        }

        headers := jws.NewHeaders()
        for key, val := range c.protected {
                if err := headers.Set(key, val); err != nil {
                        return nil, err
                }
        }

        cicToken, err := jws.Sign(
                payloadDecoded,
                jws.WithKey(
                        algorithm,
                        signer,
                        jws.WithProtectedHeaders(headers),
                ),
        )
        if err != nil {
                return nil, err
        }

        return cicToken, nil
}

func generateRand() (string, error) {
        bits := 256
        rBytes := make([]byte, bits/8)
        _, err := rand.Read(rBytes)
        if err != nil {
                return "", err
        }

        rz := hex.EncodeToString(rBytes)
        return rz, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package pktoken

import (
        "bytes"
        "fmt"

        "github.com/openpubkey/openpubkey/oidc"
)

// CompactPKToken creates a compact representation of a PK Token from a list of tokens
func CompactPKToken(tokens [][]byte, freshIDToken []byte) ([]byte, error) {
        if len(tokens) == 0 {
                return nil, fmt.Errorf("no tokens provided")
        }

        compact := [][]byte{}
        var payload []byte
        for _, tok := range tokens {
                tokProtected, tokPayload, tokSig, err := oidc.SplitCompact(tok)
                if err != nil {
                        return nil, err
                }
                if payload != nil {
                        if !bytes.Equal(payload, tokPayload) {
                                return nil, fmt.Errorf("payloads in tokens are not the same got %s and %s", payload, tokPayload)
                        }
                } else {
                        payload = tokPayload
                }
                compact = append(compact, tokProtected, tokSig)
        }
        // prepend the payload to the front
        compact = append([][]byte{payload}, compact...)
        pktCom := bytes.Join(compact, []byte(":"))

        // If we have a refreshed ID Token, append it to the compact representation using "."
        if freshIDToken != nil {
                if len(bytes.Split(freshIDToken, []byte("."))) != 3 {
                        // Compact ID Token should be reformated as Base64(protected)"."Base64(payload)"."Base64(signature)
                        return nil, fmt.Errorf("invalid refreshed ID Token")
                }
                pktCom = bytes.Join([][]byte{pktCom, freshIDToken}, []byte("."))
        }
        return pktCom, nil
}

// SplitCompactPKToken breaks a compact representation of a PK Token into its constituent tokens
func SplitCompactPKToken(pktCom []byte) ([][]byte, []byte, error) {
        tokensBytes, freshIDToken, _ := bytes.Cut(pktCom, []byte("."))
        tokensParts := bytes.Split(tokensBytes, []byte(":"))

        if freshIDToken != nil && len(bytes.Split(freshIDToken, []byte("."))) != 3 {
                // Compact ID Token should be reformated as Base64(protected)"."Base64(payload)"."Base64(signature)
                return nil, nil, fmt.Errorf("invalid refreshed ID Token")
        }

        // Compact PK Token with refreshed ID Token should have at least 3 parts and should be:
        // Base64(payload)":"Base64(protected1)":"Base64(signature1)"...":"Base64(protectedN)":"Base64(signatureN)"
        if len(tokensParts) < 3 || len(tokensParts)%2 != 1 {
                return nil, nil, fmt.Errorf("invalid number of segments, got %d", len(tokensParts))
        }
        tokens := [][]byte{}
        payload := tokensParts[0]
        for i := 1; i < len(tokensParts); i += 2 {
                // We return each token in JWT compact format (Base64(protected)"."Base64(payload)"."Base64(signature))
                token := bytes.Join([][]byte{tokensParts[i], payload, tokensParts[i+1]}, []byte("."))
                tokens = append(tokens, token)
        }
        return tokens, freshIDToken, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package pktoken

import (
        "encoding/json"
        "fmt"
)

type CosignerClaims struct {
        Issuer      string `json:"iss"`
        KeyID       string `json:"kid"`
        Algorithm   string `json:"alg"`
        AuthID      string `json:"eid"`
        AuthTime    int64  `json:"auth_time"`
        IssuedAt    int64  `json:"iat"` // may differ from auth_time because of refresh
        Expiration  int64  `json:"exp"`
        RedirectURI string `json:"ruri"`
        Nonce       string `json:"nonce"`
        Typ         string `json:"typ"`
}

func (p *PKToken) ParseCosignerClaims() (*CosignerClaims, error) {
        protected, err := json.Marshal(p.Cos.ProtectedHeaders())
        if err != nil {
                return nil, err
        }

        var claims CosignerClaims
        if err := json.Unmarshal(protected, &claims); err != nil {
                return nil, err
        }

        // Check that all fields are present
        var missing []string
        if claims.Issuer == "" {
                missing = append(missing, `iss`)
        }
        if claims.KeyID == "" {
                missing = append(missing, `kid`)
        }
        if claims.Algorithm == "" {
                missing = append(missing, `alg`)
        }
        if claims.AuthID == "" {
                missing = append(missing, `eid`)
        }
        if claims.AuthTime == 0 {
                missing = append(missing, `auth_time`)
        }
        if claims.IssuedAt == 0 {
                missing = append(missing, `iat`)
        }
        if claims.Expiration == 0 {
                missing = append(missing, `exp`)
        }
        if claims.RedirectURI == "" {
                missing = append(missing, `ruri`)
        }
        if claims.Nonce == "" {
                missing = append(missing, `nonce`)
        }

        if len(missing) > 0 {
                return nil, fmt.Errorf("cosigner protect header missing required headers: %v", missing)
        }

        return &claims, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mocks

import (
        "context"
        "crypto"
        "fmt"
        "testing"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/providers/mocks"
        "github.com/openpubkey/openpubkey/util"
        "github.com/stretchr/testify/require"
)

func GenerateMockPKToken(t *testing.T, signingKey crypto.Signer, alg jwa.KeyAlgorithm) (*pktoken.PKToken, error) {
        options := &MockPKTokenOpts{
                GQSign:         false,
                CommitType:     providers.CommitTypesEnum.NONCE_CLAIM,
                CorrectCicHash: true,
                CorrectCicSig:  true,
        }
        pkt, _, err := GenerateMockPKTokenWithOpts(t, signingKey, alg, mocks.DefaultIDTokenTemplate(), options)
        return pkt, err
}

func GenerateMockPKTokenGQ(t *testing.T, signingKey crypto.Signer, alg jwa.KeyAlgorithm) (*pktoken.PKToken, error) {
        options := &MockPKTokenOpts{
                GQSign:         true,
                CommitType:     providers.CommitTypesEnum.NONCE_CLAIM,
                CorrectCicHash: true,
                CorrectCicSig:  true,
        }
        pkt, _, err := GenerateMockPKTokenWithOpts(t, signingKey, alg, mocks.DefaultIDTokenTemplate(), options)
        return pkt, err
}

type MockPKTokenOpts struct {
        GQSign         bool
        CommitType     providers.CommitType
        GQOnly         bool
        CorrectCicHash bool
        CorrectCicSig  bool
}

func GenerateMockPKTokenWithOpts(t *testing.T, signingKey crypto.Signer, alg jwa.KeyAlgorithm,
        idtTemplate mocks.IDTokenTemplate, options *MockPKTokenOpts) (*pktoken.PKToken, *mocks.MockProviderBackend, error) {

        jwkKey, err := jwk.PublicKeyOf(signingKey)
        if err != nil {
                return nil, nil, err
        }
        err = jwkKey.Set(jwk.AlgorithmKey, alg)
        if err != nil {
                return nil, nil, err
        }

        cic, err := clientinstance.NewClaims(jwkKey, map[string]any{})
        require.NoError(t, err)

        // Set gqOnly to gqCommitment since gqCommitment requires gqOnly
        gqOnly := options.CommitType.GQCommitment

        providerOpts := providers.MockProviderOpts{
                GQSign:     options.GQSign,
                CommitType: options.CommitType,
                NumKeys:    2,
                VerifierOpts: providers.ProviderVerifierOpts{
                        SkipClientIDCheck: false,
                        GQOnly:            gqOnly,
                        CommitType:        options.CommitType,
                        ClientID:          "mockClient-ID",
                },
        }

        op, backend, _, err := providers.NewMockProvider(providerOpts)
        require.NoError(t, err)
        opSignKey, keyID, _ := backend.RandomSigningKey()
        idtTemplate.KeyID = keyID
        idtTemplate.SigningKey = opSignKey

        switch options.CommitType {
        case providers.CommitTypesEnum.NONCE_CLAIM:
                idtTemplate.CommitFunc = mocks.AddNonceCommit
        case providers.CommitTypesEnum.AUD_CLAIM:
                idtTemplate.CommitFunc = mocks.AddAudCommit
        case providers.CommitTypesEnum.GQ_BOUND:
        default:
                return nil, nil, fmt.Errorf("unknown CommitType: %v", options.CommitType)
        }

        backend.SetIDTokenTemplate(&idtTemplate)

        tokens, err := op.RequestTokens(context.Background(), cic)
        if err != nil {
                return nil, nil, err
        }
        idToken := tokens.IDToken

        // Return a PK Token where the CIC which doesn't match the commitment
        if !options.CorrectCicHash {
                // overwrite the cic with a new cic with a different hash
                cic, err = clientinstance.NewClaims(jwkKey, map[string]any{"cause": "differentCicHash"})
                if err != nil {
                        return nil, nil, err
                }
        }

        // Return a PK Token where the CIC that is signed by the wrong key
        if !options.CorrectCicSig {
                // overwrite the signkey with a new key
                signingKey, err = util.GenKeyPair(alg)
                require.NoError(t, err)

                jwkKey, err = jwk.PublicKeyOf(signingKey)
                if err != nil {
                        return nil, nil, err
                }
                err = jwkKey.Set(jwk.AlgorithmKey, alg)
                if err != nil {
                        return nil, nil, err
                }
        }

        // Sign mock id token payload with cic headers
        cicToken, err := cic.Sign(signingKey, jwkKey.Algorithm(), idToken)
        if err != nil {
                return nil, nil, err
        }

        // Combine two tokens into a PK Token
        pkt, err := pktoken.New(idToken, cicToken)
        if err != nil {
                return nil, nil, err
        }
        return pkt, backend, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package pktoken

import (
        "crypto"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jws"
)

// Options configures VerifySignedMessage behavior
type Options struct {
        Typ string // Override for the expected typ value
}

type OptionFunc func(*Options)

// WithTyp sets a custom typ value for verification
func WithTyp(typ string) OptionFunc {
        return func(o *Options) {
                o.Typ = typ
        }
}

// NewSignedMessage signs a message with the signer provided. The signed
// message is OSM (OpenPubkey Signed Message) which is a type of
// JWS (JSON Web Signature). OSMs commit to the PK Token which was used
// to generate the OSM.
func (p *PKToken) NewSignedMessage(content []byte, signer crypto.Signer) ([]byte, error) {
        cic, err := p.GetCicValues()
        if err != nil {
                return nil, err
        }

        pktHash, err := p.Hash()
        if err != nil {
                return nil, err
        }

        // Create our headers as defined by section 3.5 of the OpenPubkey paper
        protected := jws.NewHeaders()
        if err := protected.Set("alg", cic.PublicKey().Algorithm()); err != nil {
                return nil, err
        }
        if err := protected.Set("kid", pktHash); err != nil {
                return nil, err
        }
        if err := protected.Set("typ", "osm"); err != nil {
                return nil, err
        }

        return jws.Sign(
                content,
                jws.WithKey(
                        cic.PublicKey().Algorithm(),
                        signer,
                        jws.WithProtectedHeaders(protected),
                ),
        )
}

// VerifySignedMessage verifies that an OSM (OpenPubkey Signed Message) using
// the public key in this PK Token. If verification is successful,
// VerifySignedMessage returns the content of the signed message. Otherwise
// it returns an error explaining why verification failed.
//
// Note: VerifySignedMessage does not check this the PK Token is valid.
// The PK Token should always be verified first before calling
// VerifySignedMessage
func (p *PKToken) VerifySignedMessage(osm []byte, options ...OptionFunc) ([]byte, error) {
        // Default options
        opts := Options{
                Typ: "osm", // Default to "osm" for backward compatibility
        }
        // Apply provided options
        for _, opt := range options {
                opt(&opts)
        }

        cic, err := p.GetCicValues()
        if err != nil {
                return nil, err
        }

        message, err := jws.Parse(osm)
        if err != nil {
                return nil, err
        }

        // Check that our OSM headers are correct
        if len(message.Signatures()) != 1 {
                return nil, fmt.Errorf("expected only one signature on jwt, received %d", len(message.Signatures()))
        }
        protected := message.Signatures()[0].ProtectedHeaders()

        // Verify typ header matches expected value from options
        typ, ok := protected.Get("typ")
        if !ok {
                return nil, fmt.Errorf("missing required header `typ`")
        }
        if typ != opts.Typ {
                return nil, fmt.Errorf(`incorrect "typ" header, expected %q but received %s`, opts.Typ, typ)
        }

        // Verify key algorithm header matches cic
        if protected.Algorithm() != cic.PublicKey().Algorithm() {
                return nil, fmt.Errorf(`incorrect "alg" header, expected %s but received %s`, cic.PublicKey().Algorithm(), protected.Algorithm())
        }

        // Verify kid header matches hash of pktoken
        kid, ok := protected.Get("kid")
        if !ok {
                return nil, fmt.Errorf("missing required header `kid`")
        }

        pktHash, err := p.Hash()
        if err != nil {
                return nil, fmt.Errorf("unable to hash PK Token: %w", err)
        }

        if kid != string(pktHash) {
                return nil, fmt.Errorf(`incorrect "kid" header, expected %s but received %s`, pktHash, kid)
        }

        _, err = jws.Verify(osm, jws.WithKey(cic.PublicKey().Algorithm(), cic.PublicKey()))
        if err != nil {
                return nil, err
        }

        // Return the osm payload
        return message.Payload(), nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package pktoken

import (
        "bytes"
        "context"
        "crypto"
        "encoding/json"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"

        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"

        "github.com/openpubkey/openpubkey/util"

        _ "golang.org/x/crypto/sha3"
)

type SignatureType string

const (
        OIDC SignatureType = "JWT"
        CIC  SignatureType = "CIC"
        COS  SignatureType = "COS"
)

type Signature = jws.Signature

type PKToken struct {
        raw []byte // the original, raw representation of the object

        Payload []byte     // decoded payload
        Op      *Signature // Provider Signature
        Cic     *Signature // Client Signature
        Cos     *Signature // Cosigner Signature

        // We keep the tokens around as  unmarshalled values can no longer be verified
        OpToken  []byte // Base64 encoded ID Token signed by the OP
        CicToken []byte // Base64 encoded Token signed by the Client
        CosToken []byte // Base64 encoded Token signed by the Cosigner

        // FreshIDToken is the refreshed ID Token. It has a different payload from
        // other tokens and must be handled separately.
        // It is only used for POP Authentication
        FreshIDToken []byte // Base64 encoded Refreshed ID Token
}

// New creates a new PKToken from an ID Token and a CIC Token.
// It adds signatures for both tokens to the PK Token and returns the PK Token.
func New(idToken []byte, cicToken []byte) (*PKToken, error) {
        pkt := &PKToken{}

        if err := pkt.AddSignature(idToken, OIDC); err != nil {
                return nil, err
        }

        if err := pkt.AddSignature(cicToken, CIC); err != nil {
                return nil, err
        }

        return pkt, nil
}

// NewFromCompact creates a PK Token from a compact representation
func NewFromCompact(pktCom []byte) (*PKToken, error) {
        tokens, freshIDToken, err := SplitCompactPKToken(pktCom)
        if err != nil {
                return nil, err
        }
        pkt := &PKToken{}

        for _, token := range tokens {
                parsedToken, err := oidc.NewJwt(token)
                if err != nil {
                        return nil, err
                }
                typ := parsedToken.GetSignature().GetProtectedClaims().Type
                if typ == "" {
                        // missing typ claim, assuming this is from the OIDC provider and set typ=OIDC=JWT
                        // Okta is known not to set the typ parameter on their ID Tokens
                        // The JWT RFC-7519 encourages but does not require that typ be set saying about typ
                        //  "This parameter is ignored by JWT implementations; any processing of this parameter is
                        // performed by the JWT application.  If present, it is RECOMMENDED that its value be "JWT"
                        //  to indicate that this object is a JWT."
                        //  https://datatracker.ietf.org/doc/html/rfc7519#section-5.1
                        typ = string(OIDC)
                }

                sigType := SignatureType(typ)
                if err := pkt.AddSignature(token, sigType); err != nil {
                        return nil, err
                }
        }
        pkt.FreshIDToken = freshIDToken
        return pkt, nil
}

// Issuer returns the issuer (`iss`) of the ID Token in the PKToken.
// It extracts the issuer from the PKToken payload and returns it as a string.
func (p *PKToken) Issuer() (string, error) {
        var claims struct {
                Issuer string `json:"iss"`
        }
        if err := json.Unmarshal(p.Payload, &claims); err != nil {
                return "", fmt.Errorf("malformatted PK token claims: %w", err)
        }
        return claims.Issuer, nil
}

// Audience returns the audience (`aud`) of the ID Token in the PKToken.
// The audience is also known as the client ID.
func (p *PKToken) Audience() (string, error) {
        var claims struct {
                Audience string `json:"aud"`
        }
        if err := json.Unmarshal(p.Payload, &claims); err != nil {
                return "", fmt.Errorf("malformatted PK token claims: %w", err)
        }
        return claims.Audience, nil
}

// Subscriber returns the subscriber (`sub`) of the ID Token in the PKToken.
// This is a unique identifier for the user at the OpenID Provider.
func (p *PKToken) Subscriber() (string, error) {
        var claims struct {
                Subscriber string `json:"sub"`
        }
        if err := json.Unmarshal(p.Payload, &claims); err != nil {
                return "", fmt.Errorf("malformatted PK token claims: %w", err)
        }
        return claims.Subscriber, nil
}

// IdentityString string returns the three attributes that are used to uniquely identify a user
// in the OpenID Connect protocol: the subscriber, the issuer
func (p *PKToken) IdentityString() (string, error) {
        sub, err := p.Subscriber()
        if err != nil {
                return "", err
        }
        iss, err := p.Issuer()
        if err != nil {
                return "", err
        }
        return fmt.Sprintf("%s %s", sub, iss), nil
}

// Signs PK Token and then returns only the payload, header and signature as a JWT
func (p *PKToken) SignToken(
        signer crypto.Signer,
        alg jwa.KeyAlgorithm,
        protected map[string]any,
) ([]byte, error) {
        headers := jws.NewHeaders()
        for key, val := range protected {
                if err := headers.Set(key, val); err != nil {
                        return nil, fmt.Errorf("malformatted headers: %w", err)
                }
        }
        return jws.Sign(
                p.Payload,
                jws.WithKey(
                        alg,
                        signer,
                        jws.WithProtectedHeaders(headers),
                ),
        )
}

// AddSignature will add a signature to the PKToken with the specified signature type.
// It takes a token byte slice and a signature type as input, and returns an error if the signature cannot be added.
//
// To use AddSignature, first parse the token byte slice using the jws.Parse function to obtain a jws.Message object.
// You can then extract the signature from the message object using the Signatures method, and pass it to AddSignature along with the desired signature type.
//
// The function supports three signature types: OIDC, CIC, and COS.
// These signature types correspond to the JWTs in the PK Token.
// Depending on the signature type, the function will set the corresponding field in the PKToken struct (Op, Cic, or Cos) to the provided signature.
// It will also set the corresponding token field (OpToken, CicToken, or CosToken) to the provided token byte slice.
//
// If the signature type is not recognized, an error will be returned.
func (p *PKToken) AddSignature(token []byte, sigType SignatureType) error {
        message, err := jws.Parse(token)
        if err != nil {
                return err
        }

        // If there is no payload, we set the provided token's payload as current, otherwise
        // we make sure that the new payload matches current
        if p.Payload == nil {
                p.Payload = message.Payload()
        } else if !bytes.Equal(p.Payload, message.Payload()) {
                return fmt.Errorf("payload in the GQ token (%s) does not match the existing payload in the PK Token (%s)", p.Payload, message.Payload())
        }

        signature := message.Signatures()[0]

        if sigType == CIC || sigType == COS {
                protected := signature.ProtectedHeaders()
                if sigTypeFound, ok := protected.Get(jws.TypeKey); !ok {
                        return fmt.Errorf("required 'typ' claim not found in protected")
                } else if sigTypeFoundStr, ok := sigTypeFound.(string); !ok {
                        return fmt.Errorf("'typ' claim in protected must be a string but was a %T", sigTypeFound)
                } else if sigTypeFoundStr != string(sigType) {
                        return fmt.Errorf("incorrect 'typ' claim in protected, expected (%s), got (%s)", sigType, sigTypeFound)
                }
        }

        switch sigType {
        case OIDC:
                p.Op = signature
                p.OpToken = token
        case CIC:
                p.Cic = signature
                p.CicToken = token
        case COS:
                p.Cos = signature
                p.CosToken = token
        default:
                return fmt.Errorf("unrecognized signature type: %s", string(sigType))
        }
        return nil
}

func (p *PKToken) ProviderAlgorithm() (jwa.SignatureAlgorithm, bool) {
        alg, ok := p.Op.ProtectedHeaders().Get(jws.AlgorithmKey)
        if !ok {
                return "", false
        }

        return alg.(jwa.SignatureAlgorithm), true
}

func (p *PKToken) GetCicValues() (*clientinstance.Claims, error) {
        cicPH, err := p.Cic.ProtectedHeaders().AsMap(context.TODO())
        if err != nil {
                return nil, err
        }

        return clientinstance.ParseClaims(cicPH)
}

func (p *PKToken) Hash() (string, error) {
        /*
                We set the raw variable when unmarshalling from json (the only current string representation of a
                PK Token) so when we hash we use the same representation that was given for consistency. When the
                token being hashed is a new PK Token, we marshal it ourselves. This can introduce some issues based
                on how different languages format their json strings.
        */
        message := p.raw
        var err error
        if message == nil {
                message, err = json.Marshal(p)
                if err != nil {
                        return "", err
                }
        }

        hash := util.B64SHA3_256(message)
        return string(hash), nil
}

// Compact serializes a PK Token into a compact representation.
func (p *PKToken) Compact() ([]byte, error) {
        tokens := [][]byte{}
        if p.OpToken != nil {
                tokens = append(tokens, p.OpToken)
        }
        if p.CicToken != nil {
                tokens = append(tokens, p.CicToken)
        }
        if p.CosToken != nil {
                tokens = append(tokens, p.CosToken)
        }
        return CompactPKToken(tokens, p.FreshIDToken)
}

func (p *PKToken) MarshalJSON() ([]byte, error) {
        rawJws := oidc.Jws{
                Payload:    string(util.Base64EncodeForJWT(p.Payload)),
                Signatures: []oidc.Signature{},
        }
        var opPublicHeader map[string]any
        var err error
        if p.Op.PublicHeaders() != nil {
                if opPublicHeader, err = p.Op.PublicHeaders().AsMap(context.Background()); err != nil {
                        return nil, err
                }
        }
        if err = rawJws.AddSignature(p.OpToken, oidc.WithPublicHeader(opPublicHeader)); err != nil {
                return nil, err
        }
        if err = rawJws.AddSignature(p.CicToken); err != nil {
                return nil, err
        }
        if p.CosToken != nil {
                if err = rawJws.AddSignature(p.CosToken); err != nil {
                        return nil, err
                }
        }
        return json.Marshal(rawJws)
}

func (p *PKToken) UnmarshalJSON(data []byte) error {
        var rawJws oidc.Jws
        if err := json.Unmarshal(data, &rawJws); err != nil {
                return err
        }
        var parsed jws.Message
        if err := json.Unmarshal(data, &parsed); err != nil {
                return err
        }

        p.Payload = parsed.Payload() // base64 decoded

        opCount := 0
        cicCount := 0
        cosCount := 0
        for i, signature := range parsed.Signatures() {
                // for some reason the unmarshaled signatures have empty non-nil
                // public headers. set them to nil instead.
                public := signature.PublicHeaders()
                pubMap, _ := public.AsMap(context.Background())
                if len(pubMap) == 0 {
                        signature.SetPublicHeaders(nil)
                }

                protected := signature.ProtectedHeaders()
                var sigType SignatureType

                typeHeader, ok := protected.Get(jws.TypeKey)
                if ok {
                        sigTypeStr, ok := typeHeader.(string)
                        if !ok {
                                return fmt.Errorf(`provided "%s" is of wrong type, expected string`, jws.TypeKey)
                        }

                        sigType = SignatureType(sigTypeStr)
                } else {
                        // missing typ claim, assuming this is from the OIDC provider
                        sigType = OIDC
                }

                switch sigType {
                case OIDC:
                        opCount += 1
                        p.Op = signature
                        p.OpToken = []byte(rawJws.Signatures[i].Protected + "." + rawJws.Payload + "." + rawJws.Signatures[i].Signature)
                case CIC:
                        cicCount += 1
                        p.Cic = signature
                        p.CicToken = []byte(rawJws.Signatures[i].Protected + "." + rawJws.Payload + "." + rawJws.Signatures[i].Signature)
                case COS:
                        cosCount += 1
                        p.Cos = signature
                        p.CosToken = []byte(rawJws.Signatures[i].Protected + "." + rawJws.Payload + "." + rawJws.Signatures[i].Signature)
                default:
                        return fmt.Errorf("unrecognized signature type: %s", sigType)
                }
        }

        // Do some signature count verifications
        if opCount == 0 {
                return fmt.Errorf(`at least one signature of type "oidc" or "oidc_gq" is required`)
        } else if opCount > 1 {
                return fmt.Errorf(`only one signature of type "oidc" or "oidc_gq" is allowed, found %d`, opCount)
        }

        if cicCount == 0 {
                return fmt.Errorf(`at least one signature of type "cic" is required`)
        } else if cicCount > 1 {
                return fmt.Errorf(`only one signature of type "cic" is allowed, found %d`, cicCount)
        }

        if cosCount > 1 {
                return fmt.Errorf(`only one signature of type "cos" is allowed, found %d`, cosCount)
        }

        return nil
}

// DeepCopy creates a complete and independent copy of this PKToken,
func (p *PKToken) DeepCopy() (*PKToken, error) {
        pktJson, err := p.MarshalJSON()
        if err != nil {
                return nil, err
        }
        var pktCopy PKToken
        if err := json.Unmarshal(pktJson, &pktCopy); err != nil {
                return nil, err
        }
        pktCopy.FreshIDToken = p.FreshIDToken
        return &pktCopy, nil
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "fmt"
        "net/http"
        "time"

        "github.com/openpubkey/openpubkey/discover"
)

// AzureOptions is an options struct that configures how providers.AzureOp
// operates. See providers.GetDefaultAzureOpOptions for the recommended default
// values to use when interacting with Azure as the OpenIdProvider.
type AzureOptions struct {
        // ClientID is the client ID of the OIDC application. It should be the
        // expected "aud" claim in received ID tokens from the OP.
        ClientID string
        // Issuer is the OP's issuer URI for performing OIDC authorization and
        // discovery.
        Issuer string
        // Scopes is the list of scopes to send to the OP in the initial
        // authorization request.
        Scopes []string
        // PromptType is the type of prompt to use when requesting authorization from the user. Typically
        // this is set to "consent".
        PromptType string
        // AccessType is the type of access to request from the OP. Typically this is set to "offline".
        AccessType string
        // RedirectURIs is the list of authorized redirect URIs that can be
        // redirected to by the OP after the user completes the authorization code
        // flow exchange. Ensure that your OIDC application is configured to accept
        // these URIs otherwise an error may occur.
        RedirectURIs []string
        // GQSign denotes if the received ID token should be upgraded to a GQ token
        // using GQ signatures.
        GQSign bool
        // OpenBrowser denotes if the client's default browser should be opened
        // automatically when performing the OIDC authorization flow. This value
        // should typically be set to true, unless performing some headless
        // automation (e.g. integration tests) where you don't want the browser to
        // open.
        OpenBrowser bool
        // HttpClient is the http.Client to use when making queries to the OP (OIDC
        // code exchange, refresh, verification of ID token, fetch of JWKS endpoint,
        // etc.). If nil, then http.DefaultClient is used.
        HttpClient *http.Client
        // IssuedAtOffset configures the offset to add when validating the "iss" and
        // "exp" claims of received ID tokens from the OP.
        IssuedAtOffset time.Duration
        // TenantID is the GUID  of the Azure tenant/organization. Azure has a
        // different issuer URI for each tenant. Users that are not part of Azure
        // organization, which microsoft nicknames consumers have a default
        // tenant ID of "9188040d-6c67-4c5b-b112-36a304b66dad"
        // More details can be found at
        // https://learn.microsoft.com/en-us/entra/identity-platform/access-tokens
        TenantID string
}

func GetDefaultAzureOpOptions() *AzureOptions {
        defaultTenantID := "9188040d-6c67-4c5b-b112-36a304b66dad"
        return &AzureOptions{
                Issuer:   azureIssuer(defaultTenantID),
                ClientID: "096ce0a3-5e72-4da8-9c86-12924b294a01",
                // Scopes:   []string{"openid profile email"},
                Scopes:     []string{"openid profile email offline_access"}, // offline_access is required for refresh tokens
                PromptType: "consent",
                AccessType: "offline",
                RedirectURIs: []string{
                        "http://localhost:3000/login-callback",
                        "http://localhost:10001/login-callback",
                        "http://localhost:11110/login-callback",
                },
                GQSign:         false,
                OpenBrowser:    true,
                HttpClient:     nil,
                IssuedAtOffset: 1 * time.Minute,
        }
}

// NewAzureOp creates a Azure OP (OpenID Provider) using the
// default configurations options. It uses the OIDC Relying Party (Client)
// setup by the OpenPubkey project.
func NewAzureOp() BrowserOpenIdProvider {
        options := GetDefaultAzureOpOptions()
        return NewAzureOpWithOptions(options)
}

// NewAzureOpWithOptions creates a Azure OP with configuration specified
// using an options struct. This is useful if you want to use your own OIDC
// Client or override the configuration.
func NewAzureOpWithOptions(opts *AzureOptions) BrowserOpenIdProvider {
        return &AzureOp{
                StandardOp{
                        clientID:                  opts.ClientID,
                        Scopes:                    opts.Scopes,
                        PromptType:                opts.PromptType,
                        AccessType:                opts.AccessType,
                        RedirectURIs:              opts.RedirectURIs,
                        GQSign:                    opts.GQSign,
                        OpenBrowser:               opts.OpenBrowser,
                        HttpClient:                opts.HttpClient,
                        IssuedAtOffset:            opts.IssuedAtOffset,
                        issuer:                    opts.Issuer,
                        requestTokensOverrideFunc: nil,
                        publicKeyFinder: discover.PublicKeyFinder{
                                JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                        return discover.GetJwksByIssuer(ctx, issuer, opts.HttpClient)
                                },
                        },
                },
        }
}

type AzureOp = StandardOpRefreshable

var _ OpenIdProvider = (*AzureOp)(nil)
var _ BrowserOpenIdProvider = (*AzureOp)(nil)
var _ RefreshableOpenIdProvider = (*AzureOp)(nil)

func azureIssuer(tenantID string) string {
        return fmt.Sprintf("https://login.microsoftonline.com/%s/v2.0", tenantID)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "encoding/json"
        "fmt"
        "net/http"
        "net/url"

        "github.com/awnumar/memguard"
        "github.com/openpubkey/openpubkey/discover"
        simpleoidc "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
)

const githubIssuer = "https://token.actions.githubusercontent.com"

type GithubOp struct {
        issuer                    string // Change issuer to point this to a test issuer
        rawTokenRequestURL        string
        tokenRequestAuthToken     string
        publicKeyFinder           discover.PublicKeyFinder
        requestTokensOverrideFunc func(string) (*simpleoidc.Tokens, error)
}

var _ OpenIdProvider = (*GithubOp)(nil)

func NewGithubOpFromEnvironment() (*GithubOp, error) {
        tokenURL, err := getEnvVar("ACTIONS_ID_TOKEN_REQUEST_URL")
        if err != nil {
                return nil, err
        }
        token, err := getEnvVar("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
        if err != nil {
                return nil, err
        }

        return NewGithubOp(tokenURL, token), nil
}

func NewGithubOp(tokenURL string, token string) *GithubOp {
        op := &GithubOp{
                issuer:                    githubIssuer,
                rawTokenRequestURL:        tokenURL,
                tokenRequestAuthToken:     token,
                publicKeyFinder:           *discover.DefaultPubkeyFinder(),
                requestTokensOverrideFunc: nil,
        }
        return op
}

func buildTokenURL(rawTokenURL, audience string) (string, error) {
        parsedURL, err := url.Parse(rawTokenURL)
        if err != nil {
                return "", fmt.Errorf("failed to parse URL: %w", err)
        }

        if audience == "" {
                return "", fmt.Errorf("audience is required")
        }

        query := parsedURL.Query()
        query.Set("audience", audience)
        parsedURL.RawQuery = query.Encode()
        return parsedURL.String(), nil
}

func (g *GithubOp) PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error) {
        return g.publicKeyFinder.ByToken(ctx, g.issuer, token)
}

func (g *GithubOp) PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error) {
        return g.publicKeyFinder.ByKeyID(ctx, g.issuer, keyID)
}

func (g *GithubOp) requestTokens(ctx context.Context, cicHash string) (*memguard.LockedBuffer, error) {
        if g.requestTokensOverrideFunc != nil {
                tokens, err := g.requestTokensOverrideFunc(cicHash)
                if err != nil {
                        return nil, fmt.Errorf("error requesting ID Token: %w", err)
                }
                return memguard.NewBufferFromBytes(tokens.IDToken), nil
        }

        tokenURL, err := buildTokenURL(g.rawTokenRequestURL, cicHash)
        if err != nil {
                return nil, err
        }

        request, err := http.NewRequestWithContext(ctx, "GET", tokenURL, nil)
        if err != nil {
                return nil, err
        }

        request.Header.Set("Authorization", "Bearer "+g.tokenRequestAuthToken)

        var httpClient http.Client
        response, err := httpClient.Do(request)
        if err != nil {
                return nil, err
        }
        defer response.Body.Close()

        if response.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("received non-200 from jwt api: %s", http.StatusText(response.StatusCode))
        }

        rawBody, err := memguard.NewBufferFromEntireReader(response.Body)
        if err != nil {
                return nil, err
        }
        defer rawBody.Destroy()

        var jwt struct {
                Value json.RawMessage
        }
        err = json.Unmarshal(rawBody.Bytes(), &jwt)
        if err != nil {
                return nil, err
        }
        defer memguard.WipeBytes([]byte(jwt.Value))

        // json.RawMessage leaves the " (quotes) on the string. We need to remove the quotes
        return memguard.NewBufferFromBytes(jwt.Value[1 : len(jwt.Value)-1]), nil
}

func (g *GithubOp) RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error) {
        // Define our commitment as the hash of the client instance claims
        commitment, err := cic.Hash()
        if err != nil {
                return nil, fmt.Errorf("error calculating client instance claim commitment: %w", err)
        }

        // Use the commitment nonce to complete the OIDC flow and get an ID token from the provider
        idTokenLB, err := g.requestTokens(ctx, string(commitment))
        // idTokenLB is the ID Token in a memguard LockedBuffer, this is done
        // because the ID Token contains the OPs RSA signature which is a secret
        // in GQ signatures. For non-GQ signatures OPs RSA signature is considered
        // a public value.
        if err != nil {
                return nil, fmt.Errorf("error requesting ID Token: %w", err)
        }
        defer idTokenLB.Destroy()
        gqToken, err := CreateGQToken(ctx, idTokenLB.Bytes(), g)

        return &simpleoidc.Tokens{IDToken: gqToken}, err
}

func (g *GithubOp) Issuer() string {
        return g.issuer
}

func (g *GithubOp) VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error {
        vp := NewProviderVerifier(g.issuer, ProviderVerifierOpts{CommitType: CommitTypesEnum.AUD_CLAIM, GQOnly: true, SkipClientIDCheck: true})
        return vp.VerifyIDToken(ctx, idt, cic)
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "net/http"
        "time"

        "github.com/openpubkey/openpubkey/discover"
)

type GitlabOptions struct {
        // ClientID is the client ID of the OIDC application. It should be the
        // expected "aud" claim in received ID tokens from the OP.
        ClientID string
        // ClientSecret is the client secret of the OIDC application. Some OPs do
        // not require that this value is set.
        ClientSecret string
        // Issuer is the OP's issuer URI for performing OIDC authorization and
        // discovery.
        Issuer string
        // Scopes is the list of scopes to send to the OP in the initial
        // authorization request.
        Scopes []string
        // PromptType is the type of prompt to use when requesting authorization from the user. Typically
        // this is set to "consent".
        PromptType string
        // AccessType is the type of access to request from the OP. Typically this is set to "offline".
        AccessType string
        // RedirectURIs is the list of authorized redirect URIs that can be
        // redirected to by the OP after the user completes the authorization code
        // flow exchange. Ensure that your OIDC application is configured to accept
        // these URIs otherwise an error may occur.
        RedirectURIs []string
        // GQSign denotes if the received ID token should be upgraded to a GQ token
        // using GQ signatures.
        GQSign bool
        // OpenBrowser denotes if the client's default browser should be opened
        // automatically when performing the OIDC authorization flow. This value
        // should typically be set to true, unless performing some headless
        // automation (e.g. integration tests) where you don't want the browser to
        // open.
        OpenBrowser bool
        // HttpClient is the http.Client to use when making queries to the OP (OIDC
        // code exchange, refresh, verification of ID token, fetch of JWKS endpoint,
        // etc.). If nil, then http.DefaultClient is used.
        HttpClient *http.Client
        // IssuedAtOffset configures the offset to add when validating the "iss" and
        // "exp" claims of received ID tokens from the OP.
        IssuedAtOffset time.Duration
}

func GetDefaultGitlabOpOptions() *GitlabOptions {
        return &GitlabOptions{
                ClientID:   "8d8b7024572c7fd501f64374dec6bba37096783dfcd792b3988104be08cb6923",
                Issuer:     gitlabIssuer,
                Scopes:     []string{"openid email"},
                PromptType: "consent",
                AccessType: "offline",
                RedirectURIs: []string{
                        "http://localhost:3000/login-callback",
                        "http://localhost:10001/login-callback",
                        "http://localhost:11110/login-callback",
                },
                GQSign:         false,
                OpenBrowser:    true,
                HttpClient:     nil,
                IssuedAtOffset: 1 * time.Minute,
        }
}

func NewGitlabOpWithOptions(opts *GitlabOptions) BrowserOpenIdProvider {
        return &GitlabOp{
                StandardOp{
                        clientID:                  opts.ClientID,
                        Scopes:                    opts.Scopes,
                        PromptType:                opts.PromptType,
                        AccessType:                opts.AccessType,
                        RedirectURIs:              opts.RedirectURIs,
                        GQSign:                    opts.GQSign,
                        OpenBrowser:               opts.OpenBrowser,
                        HttpClient:                opts.HttpClient,
                        IssuedAtOffset:            opts.IssuedAtOffset,
                        issuer:                    opts.Issuer,
                        requestTokensOverrideFunc: nil,
                        publicKeyFinder: discover.PublicKeyFinder{
                                JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                        return discover.GetJwksByIssuer(ctx, issuer, opts.HttpClient)
                                },
                        },
                },
        }
}

type GitlabOp = StandardOpRefreshable

var _ OpenIdProvider = (*GitlabOp)(nil)
var _ BrowserOpenIdProvider = (*GitlabOp)(nil)
var _ RefreshableOpenIdProvider = (*GitlabOp)(nil)

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "fmt"

        "github.com/awnumar/memguard"
        "github.com/openpubkey/openpubkey/discover"
        simpleoidc "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
)

const gitlabIssuer = "https://gitlab.com"

type GitlabCiOp struct {
        issuer                    string // Change issuer to point this to a test issuer
        publicKeyFinder           discover.PublicKeyFinder
        tokenEnvVar               string
        requestTokensOverrideFunc func(string) (*simpleoidc.Tokens, error)
}

func NewGitlabCiOpFromEnvironmentDefault() *GitlabCiOp {
        return NewGitlabCiOpFromEnvironment("OPENPUBKEY_JWT")
}

func NewGitlabCiOpFromEnvironment(tokenEnvVar string) *GitlabCiOp {
        return NewGitlabCiOp(gitlabIssuer, tokenEnvVar)
}

func NewGitlabCiOp(issuer string, tokenEnvVar string) *GitlabCiOp {
        op := &GitlabCiOp{
                issuer:                    issuer,
                publicKeyFinder:           *discover.DefaultPubkeyFinder(),
                tokenEnvVar:               tokenEnvVar,
                requestTokensOverrideFunc: nil,
        }
        return op
}

func (g *GitlabCiOp) PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error) {
        return g.publicKeyFinder.ByToken(ctx, g.issuer, token)
}

func (g *GitlabCiOp) PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error) {
        return g.publicKeyFinder.ByKeyID(ctx, g.issuer, keyID)
}

func (g *GitlabCiOp) RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error) {
        // Define our commitment as the hash of the client instance claims
        cicHash, err := cic.Hash()
        if err != nil {
                return nil, fmt.Errorf("error calculating client instance claim commitment: %w", err)
        }

        var idToken []byte
        if g.requestTokensOverrideFunc != nil {
                noCicHashInIDToken := ""
                if tokens, err := g.requestTokensOverrideFunc(noCicHashInIDToken); err != nil {
                        return nil, fmt.Errorf("error requesting ID Token: %w", err)
                } else {
                        idToken = tokens.IDToken
                }
        } else {
                idTokenStr, err := getEnvVar(g.tokenEnvVar)
                if err != nil {
                        return nil, fmt.Errorf("error requesting ID Token: %w", err)
                }
                idToken = []byte(idTokenStr)
        }
        // idTokenLB is the ID Token in a memguard LockedBuffer, this is done
        // because the ID Token contains the OPs RSA signature which is a secret
        // in GQ signatures. For non-GQ signatures OPs RSA signature is considered
        // a public value.
        idTokenLB := memguard.NewBufferFromBytes([]byte(idToken))
        defer idTokenLB.Destroy()
        gqToken, err := CreateGQBoundToken(ctx, idTokenLB.Bytes(), g, string(cicHash))
        if err != nil {
                return nil, err
        }
        return &simpleoidc.Tokens{IDToken: []byte(gqToken)}, nil
}

func (g *GitlabCiOp) Issuer() string {
        return g.issuer
}

func (g *GitlabCiOp) VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error {
        vp := NewProviderVerifier(g.issuer,
                ProviderVerifierOpts{CommitType: CommitTypesEnum.GQ_BOUND, GQOnly: true, SkipClientIDCheck: true},
        )
        return vp.VerifyIDToken(ctx, idt, cic)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "net/http"

        "time"

        "github.com/openpubkey/openpubkey/discover"
)

const googleIssuer = "https://accounts.google.com"

// GoogleOptions is an options struct that configures how providers.GoogleOp
// operates. See providers.GetDefaultGoogleOpOptions for the recommended default
// values to use when interacting with Google as the OpenIdProvider.
type GoogleOptions struct {
        // ClientID is the client ID of the OIDC application. It should be the
        // expected "aud" claim in received ID tokens from the OP.
        ClientID string
        // ClientSecret is the client secret of the OIDC application. Some OPs do
        // not require that this value is set.
        ClientSecret string
        // Issuer is the OP's issuer URI for performing OIDC authorization and
        // discovery.
        Issuer string
        // Scopes is the list of scopes to send to the OP in the initial
        // authorization request.
        Scopes []string
        // PromptType is the type of prompt to use when requesting authorization from the user. Typically
        // this is set to "consent".
        PromptType string
        // AccessType is the type of access to request from the OP. Typically this is set to "offline".
        AccessType string
        // RedirectURIs is the list of authorized redirect URIs that can be
        // redirected to by the OP after the user completes the authorization code
        // flow exchange. Ensure that your OIDC application is configured to accept
        // these URIs otherwise an error may occur.
        RedirectURIs []string
        // GQSign denotes if the received ID token should be upgraded to a GQ token
        // using GQ signatures.
        GQSign bool
        // OpenBrowser denotes if the client's default browser should be opened
        // automatically when performing the OIDC authorization flow. This value
        // should typically be set to true, unless performing some headless
        // automation (e.g. integration tests) where you don't want the browser to
        // open.
        OpenBrowser bool
        // HttpClient is the http.Client to use when making queries to the OP (OIDC
        // code exchange, refresh, verification of ID token, fetch of JWKS endpoint,
        // etc.). If nil, then http.DefaultClient is used.
        HttpClient *http.Client
        // IssuedAtOffset configures the offset to add when validating the "iss" and
        // "exp" claims of received ID tokens from the OP.
        IssuedAtOffset time.Duration
}

func GetDefaultGoogleOpOptions() *GoogleOptions {
        return &GoogleOptions{
                Issuer:   googleIssuer,
                ClientID: "206584157355-7cbe4s640tvm7naoludob4ut1emii7sf.apps.googleusercontent.com",
                // The clientSecret was intentionally checked in. It holds no power. Do not report as a security issue
                // Google requires a ClientSecret even if this a public OIDC App
                ClientSecret: "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", // The client secret is a public value
                Scopes:       []string{"openid profile email"},
                PromptType:   "consent",
                AccessType:   "offline",
                RedirectURIs: []string{
                        "http://localhost:3000/login-callback",
                        "http://localhost:10001/login-callback",
                        "http://localhost:11110/login-callback",
                },
                GQSign:         false,
                OpenBrowser:    true,
                HttpClient:     nil,
                IssuedAtOffset: 1 * time.Minute,
        }
}

// NewGoogleOp creates a Google OP (OpenID Provider) using the
// default configurations options. It uses the OIDC Relying Party (Client)
// setup by the OpenPubkey project.
func NewGoogleOp() BrowserOpenIdProvider {
        options := GetDefaultGoogleOpOptions()
        return NewGoogleOpWithOptions(options)
}

// NewGoogleOpWithOptions creates a Google OP with configuration specified
// using an options struct. This is useful if you want to use your own OIDC
// Client or override the configuration.
func NewGoogleOpWithOptions(opts *GoogleOptions) BrowserOpenIdProvider {
        return &GoogleOp{
                StandardOp{
                        clientID:                  opts.ClientID,
                        ClientSecret:              opts.ClientSecret,
                        Scopes:                    opts.Scopes,
                        PromptType:                opts.PromptType,
                        AccessType:                opts.AccessType,
                        RedirectURIs:              opts.RedirectURIs,
                        GQSign:                    opts.GQSign,
                        OpenBrowser:               opts.OpenBrowser,
                        HttpClient:                opts.HttpClient,
                        IssuedAtOffset:            opts.IssuedAtOffset,
                        issuer:                    opts.Issuer,
                        requestTokensOverrideFunc: nil,
                        publicKeyFinder: discover.PublicKeyFinder{
                                JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                        return discover.GetJwksByIssuer(ctx, issuer, opts.HttpClient)
                                },
                        },
                },
        }
}

type GoogleOp = StandardOpRefreshable

var _ OpenIdProvider = (*GoogleOp)(nil)
var _ BrowserOpenIdProvider = (*GoogleOp)(nil)
var _ RefreshableOpenIdProvider = (*GoogleOp)(nil)

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "crypto"
        "crypto/rsa"
        "encoding/json"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/gq"
        "github.com/openpubkey/openpubkey/util"
)

func CreateGQToken(ctx context.Context, idToken []byte, op OpenIdProvider) ([]byte, error) {
        return createGQTokenAllParams(ctx, idToken, op, "", false)
}

func CreateGQBoundToken(ctx context.Context, idToken []byte, op OpenIdProvider, cicHash string) ([]byte, error) {
        return createGQTokenAllParams(ctx, idToken, op, cicHash, true)
}

func createGQTokenAllParams(ctx context.Context, idToken []byte, op OpenIdProvider, cicHash string, gqCommitment bool) ([]byte, error) {
        if cicHash != "" && !gqCommitment {
                // If gqCommitment is false, we will ignore the cicHash. This is a
                // misconfiguration, and we should fail because the caller is likely
                // expecting the cicHash to be included in the token.
                return nil, fmt.Errorf("misconfiguration, cicHash is set but gqCommitment is false, set gqCommitment to true to include cicHash in the gq signature")
        }
        headersB64, _, _, err := jws.SplitCompact(idToken)
        if err != nil {
                return nil, fmt.Errorf("error splitting compact ID Token: %w", err)
        }

        // TODO: We should create a util function for extracting headers from tokens
        headersJson, err := util.Base64DecodeForJWT(headersB64)
        if err != nil {
                return nil, fmt.Errorf("error base64 decoding ID Token headers: %w", err)
        }
        headers := jws.NewHeaders()
        err = json.Unmarshal(headersJson, &headers)
        if err != nil {
                return nil, fmt.Errorf("error unmarshalling ID Token headers: %w", err)
        }

        if headers.Algorithm() != "RS256" {
                return nil, fmt.Errorf("gq signatures require ID Token have signed with an RSA key, ID Token alg was (%s)", headers.Algorithm())
        }

        opKey, err := op.PublicKeyByToken(ctx, idToken)
        if err != nil {
                return nil, err
        }

        if opKey.Alg != "RS256" {
                return nil, fmt.Errorf("gq signatures require original provider to have signed with an RSA key, jWK.alg was (%s)", opKey.Alg)
        }

        rsaKey, ok := opKey.PublicKey.(*rsa.PublicKey)
        if !ok {
                return nil, fmt.Errorf("gq signatures require original provider to have signed with an RSA key")
        }
        jktB64, err := createJkt(rsaKey)
        if err != nil {
                return nil, err
        }

        if cicHash == "" {
                return gq.GQ256SignJWT(rsaKey, idToken, gq.WithExtraClaim("jkt", jktB64))
        } else {
                return gq.GQ256SignJWT(rsaKey, idToken, gq.WithExtraClaim("jkt", jktB64), gq.WithExtraClaim("cic", cicHash))
        }
}

func createJkt(publicKey crypto.PublicKey) (string, error) {
        jwkKey, err := jwk.PublicKeyOf(publicKey)
        if err != nil {
                return "", err
        }
        thumbprint, err := jwkKey.Thumbprint(crypto.SHA256)
        if err != nil {
                return "", err
        }
        return string(util.Base64EncodeForJWT(thumbprint)), nil
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "net/http"

        "time"

        "github.com/openpubkey/openpubkey/discover"
)

const helloIssuer = "https://issuer.hello.coop"

// HelloOptions is an options struct that configures how providers.HelloOp
// operates. See providers.GetDefaultGoogleOpOptions for the recommended default
// values to use when interacting with Google as the OpenIdProvider.
type HelloOptions struct {
        // ClientID is the client ID of the OIDC application. It should be the
        // expected "aud" claim in received ID tokens from the OP.
        ClientID string
        // Issuer is the OP's issuer URI for performing OIDC authorization and
        // discovery.
        Issuer string
        // Scopes is the list of scopes to send to the OP in the initial
        // authorization request.
        Scopes []string
        // PromptType is the type of prompt to use when requesting authorization from the user. Typically
        // this is set to "consent".
        PromptType string
        // AccessType is the type of access to request from the OP. Typically this is set to "offline".
        AccessType string
        // RedirectURIs is the list of authorized redirect URIs that can be
        // redirected to by the OP after the user completes the authorization code
        // flow exchange. Ensure that your OIDC application is configured to accept
        // these URIs otherwise an error may occur.
        RedirectURIs []string
        // GQSign denotes if the received ID token should be upgraded to a GQ token
        // using GQ signatures.
        GQSign bool
        // OpenBrowser denotes if the client's default browser should be opened
        // automatically when performing the OIDC authorization flow. This value
        // should typically be set to true, unless performing some headless
        // automation (e.g. integration tests) where you don't want the browser to
        // open.
        OpenBrowser bool
        // HttpClient is the http.Client to use when making queries to the OP (OIDC
        // code exchange, refresh, verification of ID token, fetch of JWKS endpoint,
        // etc.). If nil, then http.DefaultClient is used.
        HttpClient *http.Client
        // IssuedAtOffset configures the offset to add when validating the "iss" and
        // "exp" claims of received ID tokens from the OP.
        IssuedAtOffset time.Duration
}

func GetDefaultHelloOpOptions() *HelloOptions {
        return &HelloOptions{
                Issuer:     helloIssuer,
                ClientID:   "app_xejobTKEsDNSRd5vofKB2iay_2rN",
                Scopes:     []string{"openid profile email"},
                PromptType: "consent",
                AccessType: "offline",
                RedirectURIs: []string{
                        "http://localhost:3000/login-callback",
                        "http://localhost:10001/login-callback",
                        "http://localhost:11110/login-callback",
                },
                GQSign:         false,
                OpenBrowser:    true,
                HttpClient:     nil,
                IssuedAtOffset: 1 * time.Minute,
        }
}

// NewHelloOp creates a Google OP (OpenID Provider) using the
// default configurations options. It uses the OIDC Relying Party (Client)
// setup by the OpenPubkey project.
func NewHelloOp() BrowserOpenIdProvider {
        options := GetDefaultHelloOpOptions()
        return NewHelloOpWithOptions(options)
}

// NewHelloOpWithOptions creates a Hello OP with configuration specified
// using an options struct. This is useful if you want to use your own OIDC
// Client or override the configuration.
func NewHelloOpWithOptions(opts *HelloOptions) BrowserOpenIdProvider {
        return &HelloOp{
                clientID:                  opts.ClientID,
                Scopes:                    opts.Scopes,
                RedirectURIs:              opts.RedirectURIs,
                PromptType:                opts.PromptType,
                AccessType:                opts.AccessType,
                GQSign:                    opts.GQSign,
                OpenBrowser:               opts.OpenBrowser,
                HttpClient:                opts.HttpClient,
                IssuedAtOffset:            opts.IssuedAtOffset,
                issuer:                    opts.Issuer,
                requestTokensOverrideFunc: nil,
                publicKeyFinder: discover.PublicKeyFinder{
                        JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                return discover.GetJwksByIssuer(ctx, issuer, opts.HttpClient)
                        },
                },
        }
}

type HelloOp = StandardOp

var _ OpenIdProvider = (*HelloOp)(nil)
var _ BrowserOpenIdProvider = (*HelloOp)(nil)

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/discover"
        simpleoidc "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/providers/mocks"
)

const mockProviderIssuer = "https://accounts.example.com"

var _ OpenIdProvider = (*MockProvider)(nil)

type MockProviderOpts struct {
        Issuer     string
        Alg        string
        ClientID   string
        GQSign     bool
        NumKeys    int
        CommitType CommitType
        // We keep VerifierOpts as a variable separate to let us test failures
        // where the mock op does something which causes a verification failure
        VerifierOpts ProviderVerifierOpts
}

func DefaultMockProviderOpts() MockProviderOpts {
        clientID := "test_client_id"
        return MockProviderOpts{
                Issuer:     "https://accounts.example.com",
                Alg:        "RS256",
                ClientID:   clientID,
                GQSign:     false,
                NumKeys:    2,
                CommitType: CommitTypesEnum.NONCE_CLAIM,
                VerifierOpts: ProviderVerifierOpts{
                        CommitType:        CommitTypesEnum.NONCE_CLAIM,
                        ClientID:          clientID,
                        SkipClientIDCheck: false,
                        GQOnly:            false,
                },
        }
}

type MockProvider struct {
        options                   MockProviderOpts
        issuer                    string
        clientID                  string
        publicKeyFinder           discover.PublicKeyFinder
        requestTokensOverrideFunc func(string) (*simpleoidc.Tokens, error)
}

// NewMockProvider creates a new mock provider with a random signing key and a random key ID. It returns the provider,
// the mock backend, and the ID token template. Tests can use the mock backend to look up keys issued by the mock provider.
// Tests can use the ID token template to create ID tokens and test the provider's behavior when verifying incorrectly set ID Tokens.
func NewMockProvider(opts MockProviderOpts) (*MockProvider, *mocks.MockProviderBackend, *mocks.IDTokenTemplate, error) {
        if opts.Issuer == "" {
                opts.Issuer = mockProviderIssuer
        }
        if opts.Alg == "" {
                opts.Alg = "RS256"
        }
        mockBackend, err := mocks.NewMockProviderBackend(opts.Issuer, opts.Alg, opts.NumKeys)
        if err != nil {
                return nil, nil, nil, err
        }
        provider := &MockProvider{
                options:                   opts,
                issuer:                    mockBackend.Issuer,
                clientID:                  opts.ClientID,
                requestTokensOverrideFunc: mockBackend.RequestTokensOverrideFunc,
                publicKeyFinder:           mockBackend.PublicKeyFinder,
        }

        providerSigner, keyID, record := mockBackend.RandomSigningKey()
        commitmentFunc := mocks.NoClaimCommit
        if opts.CommitType.Claim == "nonce" {
                commitmentFunc = mocks.AddNonceCommit
        } else if opts.CommitType.Claim == "aud" {
                commitmentFunc = mocks.AddAudCommit
        }
        idTokenTemplate := &mocks.IDTokenTemplate{
                CommitFunc: commitmentFunc,
                Issuer:     provider.Issuer(),
                Nonce:      "empty",
                NoNonce:    false,
                Aud:        opts.ClientID,
                KeyID:      keyID,
                NoKeyID:    false,
                Alg:        record.Alg,
                NoAlg:      false,
                SigningKey: providerSigner,
        }
        if opts.CommitType.GQCommitment {
                idTokenTemplate.Aud = AudPrefixForGQCommitment
        }

        mockBackend.SetIDTokenTemplate(idTokenTemplate)
        return provider, mockBackend, idTokenTemplate, nil
}

func (m *MockProvider) requestTokens(_ context.Context, cicHash string) (*simpleoidc.Tokens, error) {
        return m.requestTokensOverrideFunc(cicHash)
}

func (m *MockProvider) RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error) {
        if m.options.CommitType.GQCommitment && !m.options.GQSign {
                // Catch misconfigurations in tests
                return nil, fmt.Errorf("if GQCommitment is true then GQSign must also be true")
        }

        // Define our commitment as the hash of the client instance claims
        cicHash, err := cic.Hash()
        if err != nil {
                return nil, fmt.Errorf("error calculating client instance claim commitment: %w", err)
        }
        tokens, err := m.requestTokens(ctx, string(cicHash))
        if err != nil {
                return nil, err
        }
        if m.options.CommitType.GQCommitment {
                if tokens.IDToken, err = CreateGQBoundToken(ctx, tokens.IDToken, m, string(cicHash)); err != nil {
                        return nil, err
                }
        } else if m.options.GQSign {
                if tokens.IDToken, err = CreateGQToken(ctx, tokens.IDToken, m); err != nil {
                        return nil, err
                }
        }
        return tokens, nil
}

func (m *MockProvider) RefreshTokens(ctx context.Context, _ []byte) (*simpleoidc.Tokens, error) {
        tokens, err := m.requestTokensOverrideFunc("")
        if err != nil {
                return nil, err
        }
        return tokens, nil
}

func (m *MockProvider) PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error) {
        return m.publicKeyFinder.ByToken(ctx, m.issuer, token)
}

func (m *MockProvider) PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error) {
        return m.publicKeyFinder.ByKeyID(ctx, m.issuer, keyID)
}

func (m *MockProvider) Issuer() string {
        return m.issuer
}

func (m *MockProvider) ClientID() string {
        return m.clientID
}

func (m *MockProvider) VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error {
        m.options.VerifierOpts.DiscoverPublicKey = &m.publicKeyFinder //TODO: this should be set in the constructor once we have constructors for each OP
        return NewProviderVerifier(m.Issuer(), m.options.VerifierOpts).VerifyIDToken(ctx, idt, cic)
}

func (m *MockProvider) VerifyRefreshedIDToken(ctx context.Context, origIdt []byte, reIdt []byte) error {
        if err := simpleoidc.SameIdentity(origIdt, reIdt); err != nil {
                return fmt.Errorf("refreshed ID Token is for different subject than original ID Token: %w", err)
        }
        if err := simpleoidc.RequireOlder(origIdt, reIdt); err != nil {
                return fmt.Errorf("refreshed ID Token should not be issued before original ID Token: %w", err)
        }

        pkr, err := m.publicKeyFinder.ByToken(ctx, m.Issuer(), reIdt)
        if err != nil {
                return err
        }
        alg := jwa.SignatureAlgorithm(pkr.Alg)
        if _, err := jws.Verify(reIdt, jws.WithKey(alg, pkr.PublicKey)); err != nil {
                return err
        }

        return nil
}

// Mock provider that does not support refresh
type NonRefreshableOp struct {
        op *MockProvider
}

func NewNonRefreshableOp(op *MockProvider) *NonRefreshableOp {
        return &NonRefreshableOp{op: op}
}

func (nro *NonRefreshableOp) RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error) {
        tokens, err := nro.op.RequestTokens(ctx, cic)
        return &simpleoidc.Tokens{IDToken: tokens.IDToken}, err
}
func (nro *NonRefreshableOp) PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error) {
        return nro.op.PublicKeyByKeyId(ctx, keyID)
}
func (nro *NonRefreshableOp) PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error) {
        return nro.op.PublicKeyByToken(ctx, token)
}
func (nro *NonRefreshableOp) Issuer() string {
        return nro.op.Issuer()
}
func (nro *NonRefreshableOp) VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error {
        return nro.op.VerifyIDToken(ctx, idt, cic)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mocks

import (
        "context"
        "crypto"
        "crypto/ecdsa"
        "crypto/elliptic"
        "crypto/rand"
        "crypto/rsa"
        "encoding/json"
        "fmt"
        mathrand "math/rand"

        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/openpubkey/openpubkey/discover"
        "github.com/openpubkey/openpubkey/oidc"
        "golang.org/x/exp/maps"
)

type MockProviderBackend struct {
        Issuer                string
        PublicKeyFinder       discover.PublicKeyFinder
        ProviderSigningKeySet map[string]crypto.Signer            // kid (keyId) -> signing key
        ProviderPublicKeySet  map[string]discover.PublicKeyRecord // kid (keyId) -> PublicKeyRecord
        IDTokensTemplate      *IDTokenTemplate
}

func NewMockProviderBackend(issuer string, alg string, numKeys int) (*MockProviderBackend, error) {

        var providerSigningKeySet map[string]crypto.Signer
        var providerPublicKeySet map[string]discover.PublicKeyRecord
        var err error
        if alg == "RS256" {
                if providerSigningKeySet, providerPublicKeySet, err = CreateRS256KeySet(issuer, numKeys); err != nil {
                        return nil, err
                }
        } else if alg == "ES256" {
                if providerSigningKeySet, providerPublicKeySet, err = CreateES256KeySet(issuer, numKeys); err != nil {
                        return nil, err
                }
        } else {
                return nil, fmt.Errorf("unsupported provider alg: %s", alg)
        }

        return &MockProviderBackend{
                Issuer: issuer,
                PublicKeyFinder: discover.PublicKeyFinder{
                        JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                keySet := jwk.NewSet()
                                for kid, record := range providerPublicKeySet {
                                        jwkKey, err := jwk.PublicKeyOf(record.PublicKey)
                                        if err != nil {
                                                return nil, err
                                        }
                                        if err := jwkKey.Set(jwk.AlgorithmKey, record.Alg); err != nil {
                                                return nil, err
                                        }
                                        if err := jwkKey.Set(jwk.KeyIDKey, kid); err != nil {
                                                return nil, err
                                        }

                                        // Put our jwk into a set
                                        if err := keySet.AddKey(jwkKey); err != nil {
                                                return nil, err
                                        }
                                }
                                return json.MarshalIndent(keySet, "", "  ")
                        },
                },
                ProviderSigningKeySet: providerSigningKeySet,
                ProviderPublicKeySet:  providerPublicKeySet,
        }, nil
}

func (o *MockProviderBackend) GetPublicKeyFinder() *discover.PublicKeyFinder {
        return &o.PublicKeyFinder
}

func (o *MockProviderBackend) GetProviderPublicKeySet() map[string]discover.PublicKeyRecord {
        return o.ProviderPublicKeySet
}

func (o *MockProviderBackend) GetProviderSigningKeySet() map[string]crypto.Signer {
        return o.ProviderSigningKeySet
}

func (o *MockProviderBackend) SetIDTokenTemplate(template *IDTokenTemplate) {
        o.IDTokensTemplate = template
}

func (o *MockProviderBackend) RequestTokensOverrideFunc(cicHash string) (*oidc.Tokens, error) {
        o.IDTokensTemplate.AddCommit(cicHash)
        return o.IDTokensTemplate.IssueToken()
}

func (o *MockProviderBackend) RandomSigningKey() (crypto.Signer, string, discover.PublicKeyRecord) {
        keyIDs := maps.Keys(o.GetProviderPublicKeySet())
        keyID := keyIDs[mathrand.Intn(len(keyIDs))]
        return o.GetProviderSigningKeySet()[keyID], keyID, o.GetProviderPublicKeySet()[keyID]
}

func CreateRS256KeySet(issuer string, numKeys int) (map[string]crypto.Signer, map[string]discover.PublicKeyRecord, error) {
        return CreateKeySet(issuer, "RS256", numKeys)
}

func CreateES256KeySet(issuer string, numKeys int) (map[string]crypto.Signer, map[string]discover.PublicKeyRecord, error) {
        return CreateKeySet(issuer, "ES256", numKeys)
}

func CreateKeySet(issuer string, alg string, numKeys int) (map[string]crypto.Signer, map[string]discover.PublicKeyRecord, error) {
        providerSigningKeySet := map[string]crypto.Signer{}
        providerPublicKeySet := map[string]discover.PublicKeyRecord{}

        for i := 0; i < numKeys; i++ {
                kid := fmt.Sprintf("kid-%d", i)

                var signingKey crypto.Signer
                var err error
                switch alg {
                case "ES256":
                        if signingKey, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader); err != nil {
                                return nil, nil, err
                        }
                case "RS256":
                        if signingKey, err = rsa.GenerateKey(rand.Reader, 2048); err != nil {
                                return nil, nil, err
                        }
                default:
                        return nil, nil, fmt.Errorf("unsupported alg: %s", alg)
                }

                providerSigningKeySet[string(kid)] = signingKey
                providerPublicKeySet[string(kid)] = discover.PublicKeyRecord{
                        PublicKey: signingKey.Public(),
                        Alg:       alg,
                        Issuer:    issuer,
                }
        }
        return providerSigningKeySet, providerPublicKeySet, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package mocks

import (
        "crypto"
        "crypto/ecdsa"
        "crypto/rsa"
        "encoding/json"
        "fmt"
        "time"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/oidc"
)

type CommitmentType struct {
        ClaimCommitment bool
        ClaimName       string
}

type IDTokenTemplate struct {
        CommitFunc           func(*IDTokenTemplate, string)
        Issuer               string
        Nonce                string
        NoNonce              bool
        Aud                  string
        KeyID                string
        NoKeyID              bool
        Alg                  string
        NoAlg                bool // Even if NOAlg is true, we still need Alg to be set to generate the signature
        ExtraClaims          map[string]any
        ExtraProtectedClaims map[string]any
        SigningKey           crypto.Signer // The key we will use to sign the ID Token
}

func DefaultIDTokenTemplate() IDTokenTemplate {
        return IDTokenTemplate{
                CommitFunc: AddAudCommit,
                Issuer:     "mockIssuer",
                Nonce:      "empty",
                NoNonce:    false,
                Aud:        "empty",
                KeyID:      "mockKeyID",
                NoKeyID:    false,
                Alg:        "RS256",
                NoAlg:      false,
        }
}

// AddCommit adds the commitment to the CIC to the ID Token. The
// CommitmentFunc is specified allowing custom commitment functions to be specified
func (t *IDTokenTemplate) AddCommit(cicHash string) {
        t.CommitFunc(t, cicHash)
}

// TODO: Rename to IssueTokens
func (t *IDTokenTemplate) IssueToken() (*oidc.Tokens, error) {

        headers := jws.NewHeaders()
        if !t.NoAlg {
                if err := headers.Set(jws.AlgorithmKey, t.Alg); err != nil {
                        return nil, err
                }
        }
        if !t.NoKeyID {
                if err := headers.Set(jws.KeyIDKey, t.KeyID); err != nil {
                        return nil, err
                }
        }
        if err := headers.Set(jws.TypeKey, "JWT"); err != nil {
                return nil, err
        }

        if t.ExtraProtectedClaims != nil {
                for k, v := range t.ExtraProtectedClaims {
                        if err := headers.Set(k, v); err != nil {
                                return nil, err
                        }
                }
        }

        payloadMap := map[string]any{
                "sub": "me",
                "aud": t.Aud,
                "iss": t.Issuer,
                "iat": time.Now().Unix(),
                "exp": time.Now().Add(2 * time.Hour).Unix(),
        }

        if !t.NoNonce {
                payloadMap["nonce"] = t.Nonce
        }

        if t.ExtraClaims != nil {
                for k, v := range t.ExtraClaims {
                        payloadMap[k] = v
                }
        }

        payloadBytes, err := json.Marshal(payloadMap)
        if err != nil {
                return nil, err
        }

        var providerAlg jwa.KeyAlgorithm
        if _, ok := t.SigningKey.Public().(*rsa.PublicKey); ok {
                providerAlg = jwa.RS256
        } else if _, ok := t.SigningKey.Public().(*ecdsa.PublicKey); ok {
                providerAlg = jwa.ES256
        } else {
                return nil, fmt.Errorf("unsupported public key type")
        }

        if jwa.KeyAlgorithmFrom(t.Alg) != providerAlg {
                return nil, fmt.Errorf("alg in template (%s) does not match providers signing key alg (%s)", t.Alg, providerAlg)
        }

        idToken, err := jws.Sign(
                payloadBytes,
                jws.WithKey(
                        providerAlg,
                        t.SigningKey,
                        jws.WithProtectedHeaders(headers),
                ),
        )
        if err != nil {
                return nil, err
        }
        return &oidc.Tokens{
                IDToken:      idToken,
                RefreshToken: []byte("mock-refresh-token"),
                AccessToken:  []byte("mock-access-token")}, nil
}

func AddNonceCommit(idtTemp *IDTokenTemplate, cicHash string) {
        idtTemp.Nonce = cicHash
        idtTemp.NoNonce = false
}

func AddAudCommit(idtTemp *IDTokenTemplate, cicHash string) {
        idtTemp.Aud = cicHash
}

func NoClaimCommit(idtTemp *IDTokenTemplate, cicHash string) {
        // Do nothing
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "fmt"
        "net/http"
        "os"

        "github.com/openpubkey/openpubkey/discover"
        simpleoidc "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
)

// Interface for interacting with the OP (OpenID Provider) that only returns
// an ID Token
type OpenIdProvider interface {
        RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error)
        PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error)
        PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error)
        // Returns the OpenID provider issuer as seen in ID token e.g. "https://accounts.google.com"
        Issuer() string
        VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error
}

type BrowserOpenIdProvider interface {
        OpenIdProvider
        ClientID() string
        HookHTTPSession(h http.HandlerFunc)
        ReuseBrowserWindowHook(chan string)
}

// Interface for an OpenIdProvider that returns an ID Token, Refresh Token and Access Token
type RefreshableOpenIdProvider interface {
        OpenIdProvider
        RefreshTokens(ctx context.Context, refreshToken []byte) (*simpleoidc.Tokens, error)
        VerifyRefreshedIDToken(ctx context.Context, origIdt []byte, reIdt []byte) error
}

type CommitType struct {
        Claim        string
        GQCommitment bool
}

var CommitTypesEnum = struct {
        NONCE_CLAIM CommitType
        AUD_CLAIM   CommitType
        GQ_BOUND    CommitType
}{
        NONCE_CLAIM: CommitType{Claim: "nonce", GQCommitment: false},
        AUD_CLAIM:   CommitType{Claim: "aud", GQCommitment: false},
        GQ_BOUND:    CommitType{Claim: "", GQCommitment: true}, // The commitmentClaim is bound to the ID Token using only the GQ signature
}

func getEnvVar(name string) (string, error) {
        value, ok := os.LookupEnv(name)
        if !ok {
                return "", fmt.Errorf("%q environment variable not set", name)
        }
        return value, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "crypto/ecdsa"
        "crypto/rsa"
        "encoding/json"
        "fmt"
        "strings"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/discover"
        "github.com/openpubkey/openpubkey/gq"
        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/util"
)

const AudPrefixForGQCommitment = "OPENPUBKEY-PKTOKEN:"

type DefaultProviderVerifier struct {
        issuer     string
        commitType CommitType
        options    ProviderVerifierOpts
}

type ProviderVerifierOpts struct {
        // If ClientID is specified, then verification will require that the ClientID
        // be present in the audience ("aud") claim of the PK token payload
        ClientID string
        // Describes the place where the cicHash is committed to in the the ID token.
        // For instance the nonce payload claim name where the cicHash was stored during issuance
        CommitType CommitType
        // Specifies whether to skip the Client ID check, defaults to false
        SkipClientIDCheck bool
        // Custom function for discovering public key of Provider
        DiscoverPublicKey *discover.PublicKeyFinder
        // Only allows GQ signatures, a provider signature under any other algorithm
        // is seen as an error
        GQOnly bool
}

// Creates a new ProviderVerifier with required fields
//
// issuer: Is the OpenID provider issuer as seen in ID token e.g. "https://accounts.google.com"
// commitmentClaim: the ID token payload claim name where the cicHash was stored during issuance
func NewProviderVerifier(issuer string, options ProviderVerifierOpts) *DefaultProviderVerifier {
        v := &DefaultProviderVerifier{
                issuer:     issuer,
                commitType: options.CommitType,
                options:    options,
        }

        // If no custom DiscoverPublicKey function is set, set default
        if v.options.DiscoverPublicKey == nil {
                v.options.DiscoverPublicKey = discover.DefaultPubkeyFinder()
        }

        return v
}

func (v *DefaultProviderVerifier) Issuer() string {
        return v.issuer
}

func (v *DefaultProviderVerifier) VerifyIDToken(ctx context.Context, idToken []byte, cic *clientinstance.Claims) error {
        // Sanity check that if GQCommitment is enabled then the other options
        // are set correctly for doing GQ commitment verification. The intention is
        // to catch misconfigurations early and provide meaningful error messages.
        if v.options.CommitType.GQCommitment {
                if !v.options.GQOnly {
                        return fmt.Errorf("GQCommitment requires that GQOnly is true, but GQOnly is (%t)", v.options.GQOnly)
                }
                if v.commitType.Claim != "" {
                        return fmt.Errorf("GQCommitment requires that commitmentClaim is empty but commitmentClaim is (%s)", v.commitType.Claim)
                }
                if !v.options.SkipClientIDCheck {
                        // When we bind the commitment to the ID Token using GQ Signatures,
                        // We require that the audience is prefixed with
                        // "OPENPUBKEY-PKTOKEN:". Thus, the audience can't be the client-id
                        // If you are hitting this error of set SkipClientIDCheck to true
                        return fmt.Errorf("GQCommitment requires that audience (aud) is not set to client-id")
                }
        }

        idt, err := oidc.NewJwt(idToken)
        if err != nil {
                return err
        }

        // Check whether Audience claim matches provided Client ID
        // No error is thrown if option is set to skip client ID check
        if err := verifyAudience(idt, v.options.ClientID); err != nil && !v.options.SkipClientIDCheck {
                return err
        }

        algStr := idt.GetSignature().GetProtectedClaims().Alg
        if algStr == "" {
                return fmt.Errorf("provider algorithm type missing")
        }
        alg := jwa.SignatureAlgorithm(algStr)
        if alg != gq.GQ256 && v.options.GQOnly {
                return fmt.Errorf("non-GQ signatures are not supported")
        }

        switch alg {
        case gq.GQ256:
                if err := v.verifyGQSig(ctx, idt); err != nil {
                        return fmt.Errorf("error verifying OP GQ signature on PK Token: %w", err)
                }
        case jwa.RS256:
                pubKeyRecord, err := v.providerPublicKey(ctx, idToken)
                if err != nil {
                        return fmt.Errorf("failed to get OP public key: %w", err)
                }

                // Ensure that the algorithm of public key from OpenID Provider matches the algorithm specified in the ID Token
                if _, ok := pubKeyRecord.PublicKey.(*rsa.PublicKey); !ok {
                        return fmt.Errorf("public key is not an RSA public key")
                }

                if _, err := jws.Verify(idToken, jws.WithKey(alg, pubKeyRecord.PublicKey)); err != nil {
                        return err
                }
        case jwa.ES256:
                pubKeyRecord, err := v.providerPublicKey(ctx, idToken)
                if err != nil {
                        return fmt.Errorf("failed to get OP public key: %w", err)
                }

                // Ensure that the algorithm of public key from OpenID Provider matches the algorithm specified in the ID Token
                if _, ok := pubKeyRecord.PublicKey.(*ecdsa.PublicKey); !ok {
                        return fmt.Errorf("public key is not an ECDSA public key")
                }

                if _, err := jws.Verify(idToken, jws.WithKey(alg, pubKeyRecord.PublicKey)); err != nil {
                        return err
                }
        default:
                return fmt.Errorf("unsupported signature algorithm %s", alg)
        }

        if err := v.verifyCommitment(idt, cic); err != nil {
                return err
        }

        return nil
}

// This function takes in an OIDC Provider created ID token or GQ-signed modification of one and returns
// the associated public key
func (v *DefaultProviderVerifier) providerPublicKey(ctx context.Context, idToken []byte) (*discover.PublicKeyRecord, error) {
        return v.options.DiscoverPublicKey.ByToken(ctx, v.Issuer(), idToken)
}

func (v *DefaultProviderVerifier) verifyCommitment(idt *oidc.Jwt, cic *clientinstance.Claims) error {
        var claims map[string]any
        payload, err := util.Base64DecodeForJWT([]byte(idt.GetPayload()))
        if err != nil {
                return err
        }
        if err := json.Unmarshal(payload, &claims); err != nil {
                return err
        }

        expectedCommitment, err := cic.Hash()
        if err != nil {
                return err
        }

        var commitment any
        var commitmentFound bool
        if v.options.CommitType.GQCommitment {
                aud, ok := claims["aud"]
                if !ok {
                        return fmt.Errorf("require audience claim prefix missing in PK Token's GQCommitment")
                }

                // To prevent attacks where a attacker takes someone else's ID Token
                // and turns it into a PK Token using a GQCommitment, we require that
                // all GQ commitments explicitly signal they want to be used as
                // PK Tokens. To signal this, they prefix the audience (aud)
                // claim with the string "OPENPUBKEY-PKTOKEN:".
                // We reject all GQ commitment PK Tokens that don't have this prefix
                // in the aud claim.
                if _, ok := strings.CutPrefix(aud.(string), AudPrefixForGQCommitment); !ok {
                        return fmt.Errorf("audience claim in PK Token's GQCommitment must be prefixed by (%s), got (%s) instead",
                                AudPrefixForGQCommitment, aud.(string))
                }

                // Get the commitment from the GQ signed protected header claim "cic" in the ID Token
                commitment = idt.GetSignature().GetProtectedClaims().CIC
                if commitment == "" {
                        return fmt.Errorf("missing GQ commitment")
                }
        } else {
                if v.commitType.Claim == "" {
                        return fmt.Errorf("verifier configured with empty commitment claim")
                }

                commitment, commitmentFound = claims[v.commitType.Claim]
                if !commitmentFound {
                        return fmt.Errorf("missing commitment claim %s", v.commitType.Claim)
                }
        }

        if commitment != string(expectedCommitment) {
                return fmt.Errorf("commitment claim doesn't match, got %q, expected %s", commitment, string(expectedCommitment))
        }
        return nil
}

// verifyGQSig verifies the signature of a PK token with a GQ signature. The
// parameter issuer should be the issuer of the ProviderVerifier not the
// issuer of the PK Token
func (v *DefaultProviderVerifier) verifyGQSig(ctx context.Context, idt *oidc.Jwt) error {
        algStr := idt.GetSignature().GetProtectedClaims().Alg
        if algStr == "" {
                return fmt.Errorf("missing provider algorithm header")
        }
        if algStr != gq.GQ256.String() {
                return fmt.Errorf("signature is not of type GQ")
        }

        origHeaders, err := originalTokenHeaders(idt.GetRaw())
        if err != nil {
                return fmt.Errorf("malformed ID Token headers: %w", err)
        }

        origAlg := origHeaders.Algorithm()
        if origAlg != jwa.RS256 {
                return fmt.Errorf("expected original headers to contain RS256 alg, got %s", origAlg)
        }

        if idt.GetClaims().Issuer == "" {
                return fmt.Errorf("missing issuer in payload: %s", idt.GetPayload())
        }
        if idt.GetClaims().Issuer != v.issuer {
                return fmt.Errorf("issuer of ID Token (%s) doesn't match expected issuer (%s)", idt.GetClaims().Issuer, v.issuer)
        }

        publicKeyRecord, err := v.options.DiscoverPublicKey.ByToken(ctx, v.Issuer(), idt.GetRaw())
        if err != nil {
                return fmt.Errorf("failed to get provider public key: %w", err)
        }

        rsaKey, ok := publicKeyRecord.PublicKey.(*rsa.PublicKey)
        if !ok {
                return fmt.Errorf("jwk is not an RSA key")
        }
        ok, err = gq.GQ256VerifyJWT(rsaKey, idt.GetRaw())
        if err != nil {
                return err
        }
        if !ok {
                return fmt.Errorf("error verifying OP GQ signature on PK Token (ID Token invalid)")
        }
        return nil
}

func originalTokenHeaders(token []byte) (jws.Headers, error) {
        origHeadersB64, err := gq.OriginalJWTHeaders(token)
        if err != nil {
                return nil, fmt.Errorf("malformatted PK token headers: %w", err)
        }

        origHeaders, err := util.Base64DecodeForJWT(origHeadersB64)
        if err != nil {
                return nil, fmt.Errorf("error decoding original token headers: %w", err)
        }

        headers := jws.NewHeaders()
        err = json.Unmarshal(origHeaders, &headers)
        if err != nil {
                return nil, fmt.Errorf("error parsing segment: %w", err)
        }

        return headers, nil
}

func verifyAudience(idt *oidc.Jwt, clientID string) error {

        if idt.GetClaims().Audience == "" {
                return fmt.Errorf("missing audience claim")
        }

        for _, audience := range strings.Split(idt.GetClaims().Audience, ",") {
                if audience == clientID {
                        return nil
                }
        }
        return fmt.Errorf("audience does not contain clientID %s, aud = %v", clientID, idt.GetClaims().Audience)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "context"
        "fmt"
        "net/http"

        "time"

        "github.com/google/uuid"
        "github.com/openpubkey/openpubkey/discover"
        simpleoidc "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/util"
        "github.com/sirupsen/logrus"
        "github.com/zitadel/oidc/v3/pkg/client/rp"
        "github.com/zitadel/oidc/v3/pkg/oidc"
)

// StandardOpOptions is an options struct that configures how providers.StandardOp
// operates. See providers.GetDefaultStandardOpOptions for the recommended default
// values to use when using the standardOp for a custom OpenIdProvider.
type StandardOpOptions struct {
        // ClientID is the client ID of the OIDC application. It should be the
        // expected "aud" claim in received ID tokens from the OP.
        ClientID string
        // ClientSecret is the client secret of the OIDC application. Some OPs do
        // not require that this value is set.
        ClientSecret string
        // Issuer is the OP's issuer URI for performing OIDC authorization and
        // discovery.
        Issuer string
        // Scopes is the list of scopes to send to the OP in the initial
        // authorization request.
        Scopes []string
        // PromptType is the type of prompt to use when requesting authorization from the user. Typically
        // this is set to "consent".
        PromptType string
        // AccessType is the type of access to request from the OP. Typically this is set to "offline".
        AccessType string
        // RedirectURIs is the list of authorized redirect URIs that can be
        // redirected to by the OP after the user completes the authorization code
        // flow exchange. Ensure that your OIDC application is configured to accept
        // these URIs otherwise an error may occur.
        RedirectURIs []string
        // GQSign denotes if the received ID token should be upgraded to a GQ token
        // using GQ signatures.
        GQSign bool
        // OpenBrowser denotes if the client's default browser should be opened
        // automatically when performing the OIDC authorization flow. This value
        // should typically be set to true, unless performing some headless
        // automation (e.g. integration tests) where you don't want the browser to
        // open.
        OpenBrowser bool
        // HttpClient is the http.Client to use when making queries to the OP (OIDC
        // code exchange, refresh, verification of ID token, fetch of JWKS endpoint,
        // etc.). If nil, then http.DefaultClient is used.
        HttpClient *http.Client
        // IssuedAtOffset configures the offset to add when validating the "iss" and
        // "exp" claims of received ID tokens from the OP.
        IssuedAtOffset time.Duration
}

func GetDefaultStandardOpOptions(issuer string, clientID string) *StandardOpOptions {
        return &StandardOpOptions{
                Issuer:     issuer,
                ClientID:   clientID,
                Scopes:     []string{"openid profile email"},
                PromptType: "consent",
                AccessType: "offline",
                RedirectURIs: []string{
                        "http://localhost:3000/login-callback",
                        "http://localhost:10001/login-callback",
                        "http://localhost:11110/login-callback",
                },
                GQSign:         false,
                OpenBrowser:    true,
                HttpClient:     nil,
                IssuedAtOffset: 1 * time.Minute,
        }
}

// NewStandardOpWithOptions creates a standard OP with configuration specified
// using an options struct. This is useful if you want to use your own OIDC
// Client or override the configuration.
func NewStandardOpWithOptions(opts *StandardOpOptions) BrowserOpenIdProvider {
        return &StandardOp{
                clientID:                  opts.ClientID,
                Scopes:                    opts.Scopes,
                PromptType:                opts.PromptType,
                AccessType:                opts.AccessType,
                RedirectURIs:              opts.RedirectURIs,
                GQSign:                    opts.GQSign,
                OpenBrowser:               opts.OpenBrowser,
                HttpClient:                opts.HttpClient,
                IssuedAtOffset:            opts.IssuedAtOffset,
                issuer:                    opts.Issuer,
                requestTokensOverrideFunc: nil,
                publicKeyFinder: discover.PublicKeyFinder{
                        JwksFunc: func(ctx context.Context, issuer string) ([]byte, error) {
                                return discover.GetJwksByIssuer(ctx, issuer, opts.HttpClient)
                        },
                },
        }
}

type StandardOp struct {
        clientID                  string
        ClientSecret              string
        Scopes                    []string
        PromptType                string
        AccessType                string
        RedirectURIs              []string
        GQSign                    bool
        OpenBrowser               bool
        HttpClient                *http.Client
        IssuedAtOffset            time.Duration
        issuer                    string
        server                    *http.Server
        publicKeyFinder           discover.PublicKeyFinder
        requestTokensOverrideFunc func(string) (*simpleoidc.Tokens, error)
        httpSessionHook           http.HandlerFunc
        reuseBrowserWindowHook    chan string
}

type StandardOpRefreshable struct {
        StandardOp
}

var _ OpenIdProvider = (*StandardOp)(nil)
var _ BrowserOpenIdProvider = (*StandardOp)(nil)
var _ RefreshableOpenIdProvider = (*StandardOpRefreshable)(nil)

// NewStandardOp creates a standard OP (OpenID Provider) using the
// default configuration options and returns a BrowserOpenIdProvider.
func NewStandardOp(issuer string, clientID string) BrowserOpenIdProvider {
        options := GetDefaultStandardOpOptions(issuer, clientID)
        return NewStandardOpWithOptions(options)
}

func (s *StandardOp) requestTokens(ctx context.Context, cicHash string) (*simpleoidc.Tokens, error) {
        if s.requestTokensOverrideFunc != nil {
                return s.requestTokensOverrideFunc(cicHash)
        }

        redirectURI, ln, err := FindAvailablePort(s.RedirectURIs)
        if err != nil {
                return nil, err
        }
        logrus.Infof("listening on http://%s/", ln.Addr().String())
        logrus.Info("press ctrl+c to stop")

        mux := http.NewServeMux()
        s.server = &http.Server{Handler: mux}

        cookieHandler, err := configCookieHandler()
        if err != nil {
                return nil, err
        }
        options := []rp.Option{
                rp.WithCookieHandler(cookieHandler),
                rp.WithSigningAlgsFromDiscovery(),
                rp.WithVerifierOpts(
                        rp.WithIssuedAtOffset(s.IssuedAtOffset), rp.WithNonce(
                                func(ctx context.Context) string { return cicHash })),
        }
        options = append(options, rp.WithPKCE(cookieHandler))
        if s.HttpClient != nil {
                options = append(options, rp.WithHTTPClient(s.HttpClient))
        }

        // The reason we don't set the relyingParty on the struct and reuse it,
        // is because refresh requests require a slightly different set of
        // options. For instance we want the option to check the nonce (WithNonce)
        // here, but in RefreshTokens we don't want that option set because
        // a refreshed ID token doesn't have a nonce.
        relyingParty, err := rp.NewRelyingPartyOIDC(ctx,
                s.issuer, s.clientID, s.ClientSecret, redirectURI.String(),
                s.Scopes, options...)
        if err != nil {
                return nil, fmt.Errorf("error creating provider: %w", err)
        }

        state := func() string {
                return uuid.New().String()
        }

        shutdownServer := func() {
                if err := s.server.Shutdown(ctx); err != nil {
                        logrus.Errorf("Failed to shutdown http server: %v", err)
                }
        }

        chTokens := make(chan *oidc.Tokens[*oidc.IDTokenClaims], 1)
        chErr := make(chan error, 1)

        mux.Handle("/login", rp.AuthURLHandler(state, relyingParty,
                rp.WithURLParam("nonce", cicHash),
                // Select account requires that the user click the account they want to use.
                // Results in better UX than just automatically dropping them into their
                // only signed in account.
                // See prompt parameter in OIDC spec https://openid.net/specs/openid-connect-core-1_0.html#AuthRequest
                rp.WithPromptURLParam(s.PromptType),
                rp.WithURLParam("access_type", s.AccessType)),
        )

        marshalToken := func(w http.ResponseWriter, r *http.Request, retTokens *oidc.Tokens[*oidc.IDTokenClaims], state string, rp rp.RelyingParty) {
                if err != nil {
                        http.Error(w, err.Error(), http.StatusInternalServerError)
                        chErr <- err
                        return
                }

                chTokens <- retTokens

                // If defined the OIDC client hands over control of the HTTP server session to the OpenPubkey client.
                // Useful for redirecting the user's browser window that just finished OIDC Auth flow to the
                // MFA Cosigner Auth URI.
                if s.httpSessionHook != nil {
                        s.httpSessionHook(w, r)
                        defer shutdownServer() // If no http session hook is set, we do server shutdown in RequestTokens
                } else {
                        if _, err := w.Write([]byte("You may now close this window")); err != nil {
                                logrus.Error(err)
                        }
                }
        }

        callbackPath := redirectURI.Path
        mux.Handle(callbackPath, rp.CodeExchangeHandler(marshalToken, relyingParty))

        go func() {
                err := s.server.Serve(ln)
                if err != nil && err != http.ErrServerClosed {
                        logrus.Error(err)
                }
        }()

        loginURI := fmt.Sprintf("http://localhost:%s/login", redirectURI.Port())

        // If reuseBrowserWindowHook is set, don't open a new browser window
        // instead redirect the user's existing browser window
        if s.reuseBrowserWindowHook != nil {
                s.reuseBrowserWindowHook <- loginURI
        } else if s.OpenBrowser {
                logrus.Infof("Opening browser to %s ", loginURI)
                if err := util.OpenUrl(loginURI); err != nil {
                        logrus.Errorf("Failed to open url: %v", err)
                }
        } else {
                // If s.OpenBrowser is false, tell the user what URL to open.
                // This is useful when a user wants to use a different browser than the default one.
                logrus.Infof("Open your browser to: %s ", loginURI)
        }

        // If httpSessionHook is not defined shutdown the server when done,
        // otherwise keep it open for the httpSessionHook
        // If httpSessionHook is set we handle both possible cases to ensure
        // the server is shutdown:
        // 1. We shut it down if an error occurs in the marshalToken handler
        // 2. We shut it down if the marshalToken handler completes
        if s.httpSessionHook == nil {
                defer shutdownServer()
        }
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case err := <-chErr:
                if s.httpSessionHook != nil {
                        defer shutdownServer()
                }
                return nil, err
        case retTokens := <-chTokens:
                // retTokens is a zitadel/oidc struct. We turn it into our simpler token struct
                return &simpleoidc.Tokens{
                        IDToken:      []byte(retTokens.IDToken),
                        RefreshToken: []byte(retTokens.RefreshToken),
                        AccessToken:  []byte(retTokens.AccessToken)}, nil
        }
}

func (s *StandardOp) RequestTokens(ctx context.Context, cic *clientinstance.Claims) (*simpleoidc.Tokens, error) {
        // Define our commitment as the hash of the client instance claims
        cicHash, err := cic.Hash()
        if err != nil {
                return nil, fmt.Errorf("error calculating client instance claim commitment: %w", err)
        }
        tokens, err := s.requestTokens(ctx, string(cicHash))
        if err != nil {
                return nil, err
        }
        if s.GQSign {
                idToken := tokens.IDToken
                if gqToken, err := CreateGQToken(ctx, idToken, s); err != nil {
                        return nil, err
                } else {
                        tokens.IDToken = gqToken
                        return tokens, nil
                }
        }
        return tokens, nil
}

func (s *StandardOpRefreshable) RefreshTokens(ctx context.Context, refreshToken []byte) (*simpleoidc.Tokens, error) {
        cookieHandler, err := configCookieHandler()
        if err != nil {
                return nil, err
        }
        options := []rp.Option{
                rp.WithCookieHandler(cookieHandler),
                rp.WithVerifierOpts(
                        rp.WithIssuedAtOffset(s.IssuedAtOffset),
                        rp.WithNonce(nil), // disable nonce check
                ),
        }
        options = append(options, rp.WithPKCE(cookieHandler))
        if s.HttpClient != nil {
                options = append(options, rp.WithHTTPClient(s.HttpClient))
        }

        // The redirect URI is not sent in the refresh request so we set it to an empty string.
        // According to the OIDC spec the only values send on a refresh request are:
        // client_id, client_secret, grant_type, refresh_token, and scope.
        // https://openid.net/specs/openid-connect-core-1_0.html#RefreshingAccessToken
        redirectURI := ""
        relyingParty, err := rp.NewRelyingPartyOIDC(ctx, s.issuer, s.clientID,
                s.ClientSecret, redirectURI, s.Scopes, options...)
        if err != nil {
                return nil, fmt.Errorf("failed to create RP to verify token: %w", err)
        }
        retTokens, err := rp.RefreshTokens[*oidc.IDTokenClaims](ctx, relyingParty, string(refreshToken), "", "")
        if err != nil {
                return nil, err
        }

        if retTokens.RefreshToken == "" {
                // Google does not rotate refresh tokens, the one you get at the
                // beginning is the only one you'll ever get. This may not be true
                // of OPs.
                retTokens.RefreshToken = string(refreshToken)
        }

        return &simpleoidc.Tokens{
                IDToken:      []byte(retTokens.IDToken),
                RefreshToken: []byte(retTokens.RefreshToken),
                AccessToken:  []byte(retTokens.AccessToken)}, nil
}

func (s *StandardOp) PublicKeyByToken(ctx context.Context, token []byte) (*discover.PublicKeyRecord, error) {
        return s.publicKeyFinder.ByToken(ctx, s.issuer, token)
}

func (s *StandardOp) PublicKeyByKeyId(ctx context.Context, keyID string) (*discover.PublicKeyRecord, error) {
        return s.publicKeyFinder.ByKeyID(ctx, s.issuer, keyID)
}

func (s *StandardOp) Issuer() string {
        return s.issuer
}

func (s *StandardOp) ClientID() string {
        return s.clientID
}

func (s *StandardOp) VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error {
        vp := NewProviderVerifier(
                s.issuer,
                ProviderVerifierOpts{
                        CommitType:        CommitTypesEnum.NONCE_CLAIM,
                        ClientID:          s.clientID,
                        DiscoverPublicKey: &s.publicKeyFinder,
                })
        return vp.VerifyIDToken(ctx, idt, cic)
}

func (s *StandardOpRefreshable) VerifyRefreshedIDToken(ctx context.Context, origIdt []byte, reIdt []byte) error {
        if err := simpleoidc.SameIdentity(origIdt, reIdt); err != nil {
                return fmt.Errorf("refreshed ID Token is for different subject than original ID Token: %w", err)
        }
        if err := simpleoidc.RequireOlder(origIdt, reIdt); err != nil {
                return fmt.Errorf("refreshed ID Token should not be issued before original ID Token: %w", err)
        }

        options := []rp.Option{}
        if s.HttpClient != nil {
                options = append(options, rp.WithHTTPClient(s.HttpClient))
        }
        redirectURI := ""
        relyingParty, err := rp.NewRelyingPartyOIDC(ctx, s.issuer, s.clientID,
                s.ClientSecret, redirectURI, s.Scopes, options...)
        if err != nil {
                return fmt.Errorf("failed to create RP to verify token: %w", err)
        }
        _, err = rp.VerifyIDToken[*oidc.IDTokenClaims](ctx, string(reIdt), relyingParty.IDTokenVerifier())
        return err
}

// HookHTTPSession provides a means to hook the HTTP Server session resulting
// from the OpenID Provider sending an authcode to the OIDC client by
// redirecting the user's browser with the authcode supplied in the URI.
// If this hook is set, it will be called after the receiving the authcode
// but before send an HTTP response to the user. The code which sets this hook
// can choose what HTTP response to server to the user.
//
// We use this so that we can redirect the user web browser window to
// the MFA Cosigner URI after the user finishes the OIDC Auth flow. This
// method is only available to browser based providers.
func (s *StandardOp) HookHTTPSession(h http.HandlerFunc) {
        s.httpSessionHook = h
}

// ReuseBrowserWindow is needed so that do not open more than one browser window.
// If we are using a web based OpenID Provider chooser, we have already opened one
// window on the user's browser. We should reuse that window here rather than
// opening a second browser window.
func (s *StandardOp) ReuseBrowserWindowHook(h chan string) {
        s.reuseBrowserWindowHook = h
}

// GetBrowserWindowHook ris used by testing to trigger the redirect without
// calling out the OP. This is hidden by not including in the interface.
func (s *StandardOp) TriggerBrowserWindowHook(uri string) {
        s.reuseBrowserWindowHook <- uri
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "testing"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
        "github.com/openpubkey/openpubkey/util"
        "github.com/stretchr/testify/require"
)

func GenCIC(t *testing.T) *clientinstance.Claims {
        return GenCICExtra(t, map[string]any{})
}

func GenCICExtra(t *testing.T, extraClaims map[string]any) *clientinstance.Claims {
        alg := jwa.ES256
        signer, err := util.GenKeyPair(alg)
        require.NoError(t, err)
        jwkKey, err := jwk.PublicKeyOf(signer)
        require.NoError(t, err)
        err = jwkKey.Set(jwk.AlgorithmKey, alg)
        require.NoError(t, err)
        cic, err := clientinstance.NewClaims(jwkKey, extraClaims)
        require.NoError(t, err)
        return cic
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package providers

import (
        "crypto/rand"
        "fmt"
        "io"
        "net"
        "net/url"
        "strings"

        httphelper "github.com/zitadel/oidc/v3/pkg/http"
)

// FindAvailablePort attempts to open a listener on localhost until it finds one or runs out of redirectURIs to try
func FindAvailablePort(redirectURIs []string) (*url.URL, net.Listener, error) {
        var ln net.Listener
        var lnErr error
        for _, v := range redirectURIs {
                redirectURI, err := url.Parse(v)
                if err != nil {
                        return nil, nil, fmt.Errorf("malformed redirectURI specified, redirectURI was %s", v)
                }

                if !(strings.HasPrefix(redirectURI.Host, "localhost") ||
                        strings.HasPrefix(redirectURI.Host, "127.0.0.1") ||
                        strings.HasPrefix(redirectURI.Host, "0:0:0:0:0:0:0:1") ||
                        strings.HasPrefix(redirectURI.Host, "::1")) {
                        return nil, nil, fmt.Errorf("redirectURI must be localhost, redirectURI was  %s", redirectURI.Host)
                }

                lnStr := fmt.Sprintf("localhost:%s", redirectURI.Port())
                ln, lnErr = net.Listen("tcp", lnStr)
                if lnErr == nil {
                        return redirectURI, ln, nil
                }
        }
        return nil, nil, fmt.Errorf("failed to start a listener for the callback from the OP, got %w", lnErr)
}

func configCookieHandler() (*httphelper.CookieHandler, error) {
        // I've been unable to determine a scenario in which setting a hashKey and blockKey
        // on the cookie provide protection in the localhost redirect URI case. However I
        // see no harm in setting it.
        hashKey := make([]byte, 64)
        if _, err := io.ReadFull(rand.Reader, hashKey); err != nil {
                return nil, fmt.Errorf("failed to generate random keys for cookie storage")
        }
        blockKey := make([]byte, 32)
        if _, err := io.ReadFull(rand.Reader, blockKey); err != nil {
                return nil, fmt.Errorf("failed to generate random keys for cookie storage")
        }

        // OpenPubkey uses a localhost redirect URI to receive the authcode
        // from the OP. Localhost redirects use http not https. Thus, we should
        // not set these cookies as secure-only. This should be changed if
        // OpenPubkey added support for non-localhost redirect URIs.
        // WithUnsecure() is equivalent to not setting the 'secure' attribute
        // flag in an HTTP Set-Cookie header (see https://http.dev/set-cookie#secure)
        return httphelper.NewCookieHandler(hashKey, blockKey, httphelper.WithUnsecure()), nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package util

import (
        "encoding/base64"
)

var rawURLEncoding = base64.RawURLEncoding.Strict()

func Base64EncodeForJWT(decoded []byte) []byte {
        return base64Encode(decoded, rawURLEncoding)
}

func Base64DecodeForJWT(encoded []byte) ([]byte, error) {
        return base64Decode(encoded, rawURLEncoding)
}

func base64Encode(decoded []byte, encoding *base64.Encoding) []byte {
        encoded := make([]byte, encoding.EncodedLen(len(decoded)))
        encoding.Encode(encoded, decoded)
        return encoded
}

func base64Decode(encoded []byte, encoding *base64.Encoding) ([]byte, error) {
        decoded := make([]byte, encoding.DecodedLen(len(encoded)))
        n, err := encoding.Decode(decoded, encoded)
        if err != nil {
                return nil, err
        }
        return decoded[:n], nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package util

import "bytes"

func JoinJWTSegments(segments ...[]byte) []byte {
        return JoinBytes('.', segments...)
}

func JoinBytes(sep byte, things ...[]byte) []byte {
        return bytes.Join(things, []byte{sep})
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package util

import (
        "crypto"
        "crypto/ecdsa"
        "crypto/elliptic"
        "crypto/rand"
        "crypto/rsa"
        "crypto/x509"
        "encoding/pem"
        "fmt"
        "os"

        "github.com/lestrrat-go/jwx/v2/jwa"
)

func SKToX509Bytes(sk *ecdsa.PrivateKey) ([]byte, error) {
        x509Encoded, err := x509.MarshalECPrivateKey(sk)
        if err != nil {
                return nil, err
        }
        return pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: x509Encoded}), nil

}

func WriteSKFile(fpath string, sk *ecdsa.PrivateKey) error {
        pemBytes, err := SKToX509Bytes(sk)
        if err != nil {
                return err
        }

        return os.WriteFile(fpath, pemBytes, 0600)
}

func ReadSKFile(fpath string) (*ecdsa.PrivateKey, error) {
        pemBytes, err := os.ReadFile(fpath)
        if err != nil {
                return nil, err
        }
        block, _ := pem.Decode([]byte(pemBytes))
        return x509.ParseECPrivateKey(block.Bytes)
}

func GenKeyPair(alg jwa.KeyAlgorithm) (crypto.Signer, error) {
        switch alg {
        case jwa.ES256:
                return ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
        case jwa.RS256: // RSASSA-PKCS-v1.5 using SHA-256
                return rsa.GenerateKey(rand.Reader, 2048)
        default:
                return nil, fmt.Errorf("unsupported algorithm: %s", alg.String())
        }
}

func B64SHA3_256(msg []byte) []byte {
        h := crypto.SHA3_256.New()
        h.Write(msg)
        image := h.Sum(nil)
        return Base64EncodeForJWT(image)
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package util

import (
        "os/exec"
        "runtime"
)

// https://stackoverflow.com/questions/39320371/how-start-web-server-to-open-page-in-browser-in-golang
// open opens the specified URL in the default browser of the user.
func OpenUrl(url string) error {
        var cmd string
        var args []string

        switch runtime.GOOS {
        case "windows":
                cmd = "cmd"
                args = []string{"/c", "start"}
        case "darwin":
                cmd = "open"
        default: // "linux", "freebsd", "openbsd", "netbsd"
                cmd = "xdg-open"
        }
        args = append(args, url)
        return exec.Command(cmd, args...).Start()
}

// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package verifier

import (
        "fmt"
        "time"

        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken"
)

type ExpirationPolicy struct {
        maxAge         time.Duration
        checkMaxAge    bool
        checkExpClaim  bool
        checkRefreshed bool
}

var ExpirationPolicies = struct {
        OIDC            ExpirationPolicy // This uses the OpenID Connect expiration claim
        OIDC_REFRESHED  ExpirationPolicy // This uses the OpenID Connect expiration claim on the ID Token, if that has expired. It checks the expiration on the refreshed ID Token, a.k.a., the fresh ID Token
        MAX_AGE_24HOURS ExpirationPolicy // This replaces the OpenID Connect expiration claim with OpenPubkey 24 expiration
        MAX_AGE_48HOURS ExpirationPolicy
        MAX_AGE_1WEEK   ExpirationPolicy
        NEVER_EXPIRE    ExpirationPolicy // ID Token will never expire until the OpenID Provider rotates the ID Token
}{
        OIDC:            ExpirationPolicy{maxAge: 0, checkMaxAge: false, checkExpClaim: true},
        OIDC_REFRESHED:  ExpirationPolicy{maxAge: 0, checkMaxAge: false, checkExpClaim: false, checkRefreshed: true},
        MAX_AGE_24HOURS: ExpirationPolicy{maxAge: 24 * time.Hour, checkMaxAge: true, checkExpClaim: false},
        MAX_AGE_48HOURS: ExpirationPolicy{maxAge: 2 * 24 * time.Hour, checkMaxAge: true, checkExpClaim: false},
        MAX_AGE_1WEEK:   ExpirationPolicy{maxAge: 7 * 24 * time.Hour, checkMaxAge: true, checkExpClaim: false},
        NEVER_EXPIRE:    ExpirationPolicy{maxAge: 0, checkMaxAge: false, checkExpClaim: false},
}

// CheckExpiration checks the expiration of the PK Token against the expiration
// policy.
func (ep ExpirationPolicy) CheckExpiration(pkt *pktoken.PKToken) error {
        idt, err := oidc.NewJwt(pkt.OpToken)
        if err != nil {
                return err
        }
        idtClaims := idt.GetClaims()

        if ep.checkExpClaim {
                _, err := verifyNotExpired(idtClaims.Expiration)
                if err != nil {
                        return err
                }
        }
        if ep.checkRefreshed {
                expired, err := verifyNotExpired(idtClaims.Expiration)

                // If the id token is expired, verify against the refreshed id token
                if expired {
                        if pkt.FreshIDToken == nil {
                                return fmt.Errorf("ID token is expired and no refresh token found")
                        }
                        freshIdt, err := oidc.NewJwt(pkt.FreshIDToken)
                        if err != nil {
                                return err
                        }
                        _, err = verifyNotExpired(freshIdt.GetClaims().Expiration)
                        if err != nil {
                                return err
                        }
                } else if err != nil { // an non-expiration error occurred
                        return err
                }
        }

        if ep.checkMaxAge {
                _, err := checkMaxAge(idtClaims.IssuedAt, int64(ep.maxAge.Seconds()))
                if err != nil {
                        return err
                }
        }
        return nil
}

// verifyNotExpired checks the expiration of the ID Token using the exp claim.
// If expired, returns true and set an error. If an error prevents checking
// expiration it return false and the error.
func verifyNotExpired(expiration int64) (bool, error) {
        if expiration == 0 {
                return false, fmt.Errorf("missing expiration claim")
        }
        if expiration < 0 {
                return false, fmt.Errorf("expiration must be must be greater than zero (issuedAt = %v)", expiration)
        }

        // JWT expiration is "Seconds Since the Epoch"
        // RFC-7519 -Section 2 https://www.rfc-editor.org/rfc/rfc7519#section-2
        expirationTime := time.Unix(expiration, 0)
        if time.Now().After(expirationTime) {
                return true, fmt.Errorf("the ID token has expired (exp = %v)", expiration)
        }
        return false, nil
}

// checkMaxAge checks the max age of the ID Token using the issuedAt claim.
// If expired, returns true and set an error. If an error prevents checking
// expiration it return false and the error.
func checkMaxAge(issuedAt int64, maxAge int64) (bool, error) {
        if issuedAt == 0 {
                return false, fmt.Errorf("missing issuedAt claim")
        }
        if issuedAt < 0 {
                return false, fmt.Errorf("issuedAt must be must be greater than zero (issuedAt = %v)", issuedAt)
        }
        if !(maxAge > 0) {
                return false, fmt.Errorf("maxAge configuration must be greater than zero (maxAge = %v)", maxAge)
        }
        // Ensure we throw an error is something goes wrong and we get parameters so large they overflow
        if (issuedAt + maxAge) < issuedAt {
                return false, fmt.Errorf("invalid values (issuedAt = %v, maxAge = %v)", issuedAt, maxAge)
        }
        expirationTime := time.Unix(issuedAt+maxAge, 0)
        if time.Now().After(expirationTime) {
                return true, fmt.Errorf("the PK token has expired based on maxAge (issuedAt = %v, maxAge = %v, expiratedAt = %v)", issuedAt, maxAge, expirationTime)
        }
        return false, nil
}

// Copyright 2024 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package verifier

import (
        "context"
        "fmt"

        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/cosigner"
        "github.com/openpubkey/openpubkey/gq"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/pktoken/clientinstance"
)

type ProviderVerifier interface {
        // Returns the OpenID provider issuer as seen in ID token e.g. "https://accounts.google.com"
        Issuer() string
        VerifyIDToken(ctx context.Context, idt []byte, cic *clientinstance.Claims) error
}

type ProviderVerifierExpires struct {
        ProviderVerifier
        Expiration ExpirationPolicy
}

func (p ProviderVerifierExpires) ExpirationPolicy() ExpirationPolicy {
        return p.Expiration
}

type RefreshableProviderVerifier interface {
        VerifyRefreshedIDToken(ctx context.Context, origIdt []byte, reIdt []byte) error
}

type CosignerVerifier interface {
        Issuer() string
        Strict() bool // Whether or not a given cosigner MUST be present for successful verification
        VerifyCosigner(ctx context.Context, pkt *pktoken.PKToken) error
}

type VerifierOpts func(*Verifier) error

// RequireRefreshedIDToken instructs the verifier to check that
// an unexpired, refreshed ID token is set on the PKToken.
func RequireRefreshedIDToken() VerifierOpts {
        return func(v *Verifier) error {
                v.requireRefreshedIDToken = true
                return nil
        }
}

func WithExpirationPolicy(expirationPolicy ExpirationPolicy) VerifierOpts {
        return func(v *Verifier) error {
                v.defaultExpirationPolicy = &expirationPolicy
                return nil
        }
}

func WithCosignerVerifiers(verifiers ...*cosigner.DefaultCosignerVerifier) VerifierOpts {
        return func(v *Verifier) error {
                for _, verifier := range verifiers {
                        if _, ok := v.cosigners[verifier.Issuer()]; ok {
                                return fmt.Errorf("cosigner verifier found with duplicate issuer: %s", verifier.Issuer())
                        }
                        v.cosigners[verifier.Issuer()] = verifier
                }
                return nil
        }
}

type Check func(*Verifier, *pktoken.PKToken) error

func GQOnly() Check {
        return func(_ *Verifier, pkt *pktoken.PKToken) error {
                alg, ok := pkt.ProviderAlgorithm()
                if !ok {
                        return fmt.Errorf("missing provider algorithm header")
                }

                if alg != gq.GQ256 {
                        return fmt.Errorf("non-GQ signatures are not supported")
                }
                return nil
        }
}

type Verifier struct {
        providers map[string]ProviderVerifier
        cosigners map[string]CosignerVerifier
        // Sets the default expiration policy to use
        defaultExpirationPolicy *ExpirationPolicy
        requireRefreshedIDToken bool
}

func New(verifier ProviderVerifier, options ...VerifierOpts) (*Verifier, error) {
        return NewFromMany([]ProviderVerifier{verifier}, options...)
}

func NewFromMany(verifiers []ProviderVerifier, options ...VerifierOpts) (*Verifier, error) {
        v := &Verifier{
                providers: map[string]ProviderVerifier{},
                cosigners: map[string]CosignerVerifier{},
                // For user access we override the ID Token expiration claim
                // and instead have tokens expire after 24 hours so that
                // users don't have log back in every hour.
                defaultExpirationPolicy: &ExpirationPolicies.MAX_AGE_24HOURS,
        }

        for _, verifier := range verifiers {
                if _, ok := v.providers[verifier.Issuer()]; ok {
                        return nil, fmt.Errorf("provider verifier found with duplicate issuer: %s", verifier.Issuer())
                }
                v.providers[verifier.Issuer()] = verifier
        }

        for _, option := range options {
                if err := option(v); err != nil {
                        return nil, err
                }
        }

        if v.defaultExpirationPolicy == nil {
                // Default to 24 hours if no expiration policy is set
                v.defaultExpirationPolicy = &ExpirationPolicies.MAX_AGE_24HOURS
        }

        return v, nil
}

// Verifies whether a PK token is valid and matches all expected claims.
//
// extraChecks: Allows for optional specification of additional checks
func (v *Verifier) VerifyPKToken(
        ctx context.Context,
        pkt *pktoken.PKToken,
        extraChecks ...Check,
) error {
        // Don't even bother doing anything if the user's isn't valid
        if err := verifyCicSignature(pkt); err != nil {
                return fmt.Errorf("error verifying client signature on PK Token: %w", err)
        }

        issuer, err := pkt.Issuer()
        if err != nil {
                return err
        }

        providerVerifier, ok := v.providers[issuer]
        if !ok {
                var knownIssuers []string
                for k := range v.providers {
                        knownIssuers = append(knownIssuers, k)
                }
                return fmt.Errorf("unrecognized issuer: %s, issuers known: %v", issuer, knownIssuers)
        }

        cic, err := pkt.GetCicValues()
        if err != nil {
                return err
        }
        if err := providerVerifier.VerifyIDToken(ctx, pkt.OpToken, cic); err != nil {
                return err
        }

        // If expiration has been set for this provider verifier use it to check expiration
        if providerVerifierExpires, ok := providerVerifier.(ProviderVerifierExpires); ok {
                if err := providerVerifierExpires.ExpirationPolicy().CheckExpiration(pkt); err != nil {
                        return err
                }
        } else if err := v.defaultExpirationPolicy.CheckExpiration(pkt); err != nil {
                // Otherwise use the default expiration policy
                return err
        }

        if v.requireRefreshedIDToken {
                if reProviderVerifier, ok := providerVerifier.(RefreshableProviderVerifier); !ok {
                        return fmt.Errorf("refreshed ID Token verification required but provider verifier (issuer=%s) does not support it", issuer)
                } else {
                        if pkt.FreshIDToken == nil {
                                return fmt.Errorf("no refreshed ID Token set")
                        }
                        if err := reProviderVerifier.VerifyRefreshedIDToken(ctx, pkt.OpToken, pkt.FreshIDToken); err != nil {
                                return err
                        }
                }
        }

        if len(v.cosigners) > 0 {
                if pkt.Cos == nil {
                        // If there's no cosigner signature and any provided cosigner verifiers are strict, then return error
                        for _, cosignerVerifier := range v.cosigners {
                                if cosignerVerifier.Strict() {
                                        return fmt.Errorf("missing required cosigner signature by %s", cosignerVerifier.Issuer())
                                }
                        }
                } else {
                        cosignerClaims, err := pkt.ParseCosignerClaims()
                        if err != nil {
                                return err
                        }

                        cosignerVerifier, ok := v.cosigners[cosignerClaims.Issuer]
                        if !ok {
                                // If other cosigners are present, do we accept?
                                return fmt.Errorf("unrecognized cosigner %s", cosignerClaims.Issuer)
                        }

                        // Verify cosigner signature
                        if err := cosignerVerifier.VerifyCosigner(ctx, pkt); err != nil {
                                return err
                        }

                        // If any other cosigner verifiers are set to strict but aren't present, then return error
                        for _, cosignerVerifier := range v.cosigners {
                                if cosignerVerifier.Strict() && cosignerVerifier.Issuer() != cosignerClaims.Issuer {
                                        return fmt.Errorf("missing required cosigner signature by %s", cosignerVerifier.Issuer())
                                }
                        }
                }
        }
        // Cycles through any provided additional checks and returns the first error, if any.
        for _, check := range extraChecks {
                if err := check(v, pkt); err != nil {
                        return err
                }
        }

        return nil
}

func verifyCicSignature(pkt *pktoken.PKToken) error {
        cic, err := pkt.GetCicValues()
        if err != nil {
                return err
        }

        _, err = jws.Verify(pkt.CicToken, jws.WithKey(cic.PublicKey().Algorithm(), cic.PublicKey()))
        return err
}