// 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 commands

import (
        "errors"
        "fmt"
        "os"

        "github.com/openpubkey/opkssh/policy"
)

// AddCmd provides functionality to read and update the opkssh policy file
type AddCmd struct {
        HomePolicyLoader   *policy.HomePolicyLoader
        SystemPolicyLoader *policy.SystemPolicyLoader

        // Username is the username to lookup when the system policy file cannot be
        // read and we fallback to the user's policy file.
        //
        // See AddCmd.LoadPolicy for more details.
        Username string
}

// LoadPolicy reads the opkssh policy at the policy.SystemDefaultPolicyPath. If
// there is a permission error when reading this file, then the user's local
// policy file (defined as ~/.opk/auth_id where ~ maps to AddCmd.Username's
// home directory) is read instead.
//
// If successful, returns the parsed policy and filepath used to read the
// policy. Otherwise, a non-nil error is returned.
func (a *AddCmd) LoadPolicy() (*policy.Policy, string, error) {
        // Try to read system policy first
        systemPolicy, _, err := a.SystemPolicyLoader.LoadSystemPolicy()
        if err != nil {
                if errors.Is(err, os.ErrPermission) {
                        // If current process doesn't have permission, try reading the user
                        // policy file.
                        userPolicy, policyFilePath, err := a.HomePolicyLoader.LoadHomePolicy(a.Username, false)
                        if err != nil {
                                return nil, "", err
                        }
                        return userPolicy, policyFilePath, nil
                } else {
                        // Non-permission error (e.g. system policy file missing or invalid
                        // permission bits set). Return error
                        return nil, "", err
                }
        }

        return systemPolicy, policy.SystemDefaultPolicyPath, nil
}

// GetPolicyPath returns the path to the policy file that the current command
// will write to and a boolean to flag the path is for home policy.
// True means home policy, false means system policy.
func (a *AddCmd) GetPolicyPath(principal string, userEmail string, issuer string) (string, bool, error) {
        // Try to read system policy first
        _, _, err := a.SystemPolicyLoader.LoadSystemPolicy()
        if err != nil {
                if errors.Is(err, os.ErrPermission) {
                        // If current process doesn't have permission, try reading the user
                        // policy file.
                        policyFilePath, err := a.HomePolicyLoader.UserPolicyPath(a.Username)
                        if err != nil {
                                return "", false, err
                        }
                        return policyFilePath, false, nil
                } else {
                        // Non-permission error (e.g. system policy file missing or invalid
                        // permission bits set). Return error
                        return "", false, err
                }
        }
        return policy.SystemDefaultPolicyPath, true, nil
}

// Run adds a new allowed principal to the user whose email is equal to
// userEmail. The policy file is read and modified.
//
// If successful, returns the policy filepath updated. Otherwise, returns a
// non-nil error
func (a *AddCmd) Run(principal string, userEmail string, issuer string) (string, error) {
        policyPath, useSystemPolicy, err := a.GetPolicyPath(principal, userEmail, issuer)
        if err != nil {
                return "", fmt.Errorf("failed to load policy: %w", err)
        }

        var policyLoader *policy.PolicyLoader
        if useSystemPolicy {
                policyLoader = a.SystemPolicyLoader.PolicyLoader
        } else {
                policyLoader = a.HomePolicyLoader.PolicyLoader
        }

        err = policyLoader.CreateIfDoesNotExist(policyPath)
        if err != nil {
                return "", fmt.Errorf("failed to create policy file: %w", err)
        }

        // Read current policy
        currentPolicy, policyFilePath, err := a.LoadPolicy()
        if err != nil {
                return "", fmt.Errorf("failed to load current policy: %w", err)
        }

        // Update policy
        currentPolicy.AddAllowedPrincipal(principal, userEmail, issuer)

        // Dump contents back to disk
        err = policyLoader.Dump(currentPolicy, policyFilePath)
        if err != nil {
                return "", fmt.Errorf("failed to write updated policy: %w", err)
        }

        return policyFilePath, 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 commands

import (
        "context"
        "crypto"
        "encoding/base64"
        "encoding/json"
        "encoding/pem"
        "errors"
        "fmt"
        "io"
        "log"
        "os"
        "path/filepath"
        "strings"
        "time"

        "github.com/lestrrat-go/jwx/v2/jwa"
        "github.com/lestrrat-go/jwx/v2/jws"
        "github.com/openpubkey/openpubkey/client"
        "github.com/openpubkey/openpubkey/client/choosers"
        "github.com/openpubkey/openpubkey/oidc"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/util"
        "github.com/openpubkey/opkssh/sshcert"
        "golang.org/x/crypto/ssh"
)

type LoginCmd struct {
        autoRefresh           bool
        logDir                string
        disableBrowserOpenArg bool
        printIdTokenArg       bool
        providerArg           string
        providerFromLdFlags   providers.OpenIdProvider
        pkt                   *pktoken.PKToken
        signer                crypto.Signer
        alg                   jwa.SignatureAlgorithm
        client                *client.OpkClient
        principals            []string
}

func NewLogin(autoRefresh bool, logDir string, disableBrowserOpenArg bool, printIdTokenArg bool, providerArg string, providerFromLdFlags providers.OpenIdProvider) *LoginCmd {
        return &LoginCmd{
                autoRefresh:           autoRefresh,
                logDir:                logDir,
                disableBrowserOpenArg: disableBrowserOpenArg,
                printIdTokenArg:       printIdTokenArg,
                providerArg:           providerArg,
                providerFromLdFlags:   providerFromLdFlags,
        }
}

func (l *LoginCmd) Run(ctx context.Context) error {
        // If a log directory was provided, write any logs to a file in that directory AND stdout
        if l.logDir != "" {
                logFilePath := filepath.Join(l.logDir, "opkssh.log")
                logFile, err := os.OpenFile(logFilePath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0660)
                if err != nil {
                        log.Printf("Failed to open log for writing: %v \n", err)
                }
                defer logFile.Close()
                multiWriter := io.MultiWriter(os.Stdout, logFile)
                log.SetOutput(multiWriter)
        } else {
                log.SetOutput(os.Stdout)
        }

        openBrowser := !l.disableBrowserOpenArg

        // If the user has supplied commandline arguments for the provider, use those instead of the web chooser
        var provider providers.OpenIdProvider
        if l.providerArg != "" {
                parts := strings.Split(l.providerArg, ",")
                if len(parts) != 2 && len(parts) != 3 {
                        return fmt.Errorf("invalid provider argument format. Expected format <issuer>,<client_id> or <issuer>,<client_id>,<client_secret> got (%s)", l.providerArg)
                }
                issuerArg := parts[0]
                clientIDArg := parts[1]

                if !strings.HasPrefix(issuerArg, "https://") {
                        return fmt.Errorf("invalid provider issuer value. Expected issuer to start with 'https://' got (%s) \n", issuerArg)
                }

                if clientIDArg == "" {
                        return fmt.Errorf("invalid provider client-ID value got (%s) \n", clientIDArg)
                }

                if strings.HasPrefix(issuerArg, "https://accounts.google.com") {
                        // The Google OP is strange in that it requires a client secret even if this is a public OIDC App.
                        // Despite its name the Google OP client secret is a public value.
                        if len(parts) != 3 {
                                return fmt.Errorf("invalid provider argument format. Expected format for google: <issuer>,<client_id>,<client_secret> got (%s)", l.providerArg)
                        }
                        clientSecretArg := parts[2]
                        if clientSecretArg == "" {
                                return fmt.Errorf("invalid provider client secret value got (%s) \n", clientSecretArg)
                        }

                        opts := providers.GetDefaultGoogleOpOptions()
                        opts.Issuer = issuerArg
                        opts.ClientID = clientIDArg
                        opts.ClientSecret = clientSecretArg
                        opts.GQSign = false
                        opts.OpenBrowser = openBrowser
                        provider = providers.NewGoogleOpWithOptions(opts)
                } else if strings.HasPrefix(issuerArg, "https://login.microsoftonline.com") {
                        opts := providers.GetDefaultAzureOpOptions()
                        opts.Issuer = issuerArg
                        opts.ClientID = clientIDArg
                        opts.GQSign = false
                        opts.OpenBrowser = openBrowser
                        provider = providers.NewAzureOpWithOptions(opts)
                } else if strings.HasPrefix(issuerArg, "https://gitlab.com") {
                        opts := providers.GetDefaultGitlabOpOptions()
                        opts.Issuer = issuerArg
                        opts.ClientID = clientIDArg
                        opts.GQSign = false
                        opts.OpenBrowser = openBrowser
                        provider = providers.NewGitlabOpWithOptions(opts)
                } else {
                        // Generic provider - Need signing, no encryption
                        opts := providers.GetDefaultGoogleOpOptions()
                        opts.Issuer = issuerArg
                        opts.ClientID = clientIDArg
                        opts.ClientSecret = "" // No client secret for generic providers unless specified
                        opts.GQSign = false
                        opts.OpenBrowser = openBrowser

                        if len(parts) == 3 {
                                opts.ClientSecret = parts[2]
                        }

                        provider = providers.NewGoogleOpWithOptions(opts)
                }
        } else if l.providerFromLdFlags != nil {
                provider = l.providerFromLdFlags
        } else {
                googleOpOptions := providers.GetDefaultGoogleOpOptions()
                googleOpOptions.OpenBrowser = openBrowser
                googleOpOptions.GQSign = false
                googleOp := providers.NewGoogleOpWithOptions(googleOpOptions)

                azureOpOptions := providers.GetDefaultAzureOpOptions()
                azureOpOptions.OpenBrowser = openBrowser
                azureOpOptions.GQSign = false
                azureOp := providers.NewAzureOpWithOptions(azureOpOptions)

                gitlabOpOptions := providers.GetDefaultGitlabOpOptions()
                gitlabOpOptions.OpenBrowser = openBrowser
                gitlabOpOptions.GQSign = false
                gitlabOp := providers.NewGitlabOpWithOptions(gitlabOpOptions)

                var err error
                provider, err = choosers.NewWebChooser(
                        []providers.BrowserOpenIdProvider{googleOp, azureOp, gitlabOp},
                        !l.disableBrowserOpenArg,
                ).ChooseOp(ctx)
                if err != nil {
                        return fmt.Errorf("error selecting OpenID provider: %w", err)
                }
        }

        // Execute login command
        if l.autoRefresh {
                if providerRefreshable, ok := provider.(providers.RefreshableOpenIdProvider); ok {
                        err := LoginWithRefresh(ctx, providerRefreshable, l.printIdTokenArg)
                        if err != nil {
                                return fmt.Errorf("error logging in: %w", err)
                        }
                } else {
                        return fmt.Errorf("supplied OpenID Provider (%v) does not support auto-refresh and auto-refresh argument set to true", provider.Issuer())
                }
        } else {
                err := Login(ctx, provider, l.printIdTokenArg)
                if err != nil {
                        return fmt.Errorf("error logging in: %w", err)
                }
        }
        return nil
}

func login(ctx context.Context, provider client.OpenIdProvider, printIdToken bool) (*LoginCmd, error) {
        var err error
        alg := jwa.ES256
        signer, err := util.GenKeyPair(alg)
        if err != nil {
                return nil, fmt.Errorf("failed to generate keypair: %w", err)
        }

        opkClient, err := client.New(provider, client.WithSigner(signer, alg))
        if err != nil {
                return nil, err
        }

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

        // If principals is empty the server does not enforce any principal. The OPK
        // verifier should use policy to make this decision.
        principals := []string{}
        certBytes, seckeySshPem, err := createSSHCert(pkt, signer, principals)
        if err != nil {
                return nil, fmt.Errorf("failed to generate SSH cert: %w", err)
        }

        // Write ssh secret key and public key to filesystem
        if err := writeKeysToSSHDir(seckeySshPem, certBytes); err != nil {
                return nil, fmt.Errorf("failed to write SSH keys to filesystem: %w", err)
        }

        if printIdToken {
                idTokenStr, err := PrettyIdToken(*pkt)

                if err != nil {
                        return nil, fmt.Errorf("failed to format ID Token: %w", err)
                }

                fmt.Printf("id_token:\n%s\n", idTokenStr)
        }

        idStr, err := IdentityString(*pkt)
        if err != nil {
                return nil, fmt.Errorf("failed to parse ID Token: %w", err)
        }
        fmt.Printf("Keys generated for identity\n%s\n", idStr)

        return &LoginCmd{
                pkt:        pkt,
                signer:     signer,
                client:     opkClient,
                alg:        alg,
                principals: principals,
        }, nil
}

// Login performs the OIDC login procedure and creates the SSH certs/keys in the
// default SSH key location.
func Login(ctx context.Context, provider client.OpenIdProvider, printIdToken bool) error {
        _, err := login(ctx, provider, printIdToken)
        return err
}

// LoginWithRefresh performs the OIDC login procedure, creates the SSH
// certs/keys in the default SSH key location, and continues to run and refresh
// the PKT (and create new SSH certs) indefinitely as its token expires. This
// function only returns if it encounters an error or if the supplied context is
// cancelled.
func LoginWithRefresh(ctx context.Context, provider providers.RefreshableOpenIdProvider, printIdToken bool) error {
        if loginResult, err := login(ctx, provider, printIdToken); err != nil {
                return err
        } else {
                var claims struct {
                        Expiration int64 `json:"exp"`
                }
                if err := json.Unmarshal(loginResult.pkt.Payload, &claims); err != nil {
                        return err
                }

                for {
                        // Sleep until a minute before expiration to give us time to refresh
                        // the token and minimize any interruptions
                        untilExpired := time.Until(time.Unix(claims.Expiration, 0)) - time.Minute
                        log.Printf("Waiting for %v before attempting to refresh id_token...", untilExpired)
                        select {
                        case <-time.After(untilExpired):
                                log.Print("Refreshing id_token...")
                        case <-ctx.Done():
                                return ctx.Err()
                        }

                        refreshedPkt, err := loginResult.client.Refresh(ctx)
                        if err != nil {
                                return err
                        }
                        loginResult.pkt = refreshedPkt

                        certBytes, seckeySshPem, err := createSSHCert(loginResult.pkt, loginResult.signer, loginResult.principals)
                        if err != nil {
                                return fmt.Errorf("failed to generate SSH cert: %w", err)
                        }

                        // Write ssh secret key and public key to filesystem
                        if err := writeKeysToSSHDir(seckeySshPem, certBytes); err != nil {
                                return fmt.Errorf("failed to write SSH keys to filesystem: %w", err)
                        }

                        comPkt, err := refreshedPkt.Compact()
                        if err != nil {
                                return err
                        }

                        _, payloadB64, _, err := jws.SplitCompactString(string(comPkt))
                        if err != nil {
                                return fmt.Errorf("malformed ID token: %w", err)
                        }
                        payload, err := base64.RawURLEncoding.DecodeString(string(payloadB64))
                        if err != nil {
                                return fmt.Errorf("refreshed ID token payload is not base64 encoded: %w", err)
                        }

                        if err = json.Unmarshal(payload, &claims); err != nil {
                                return fmt.Errorf("malformed refreshed ID token payload: %w", err)
                        }
                }
        }
}

func createSSHCert(pkt *pktoken.PKToken, signer crypto.Signer, principals []string) ([]byte, []byte, error) {
        cert, err := sshcert.New(pkt, principals)
        if err != nil {
                return nil, nil, err
        }
        sshSigner, err := ssh.NewSignerFromSigner(signer)
        if err != nil {
                return nil, nil, err
        }

        signerMas, err := ssh.NewSignerWithAlgorithms(sshSigner.(ssh.AlgorithmSigner), []string{ssh.KeyAlgoECDSA256})
        if err != nil {
                return nil, nil, err
        }

        sshCert, err := cert.SignCert(signerMas)
        if err != nil {
                return nil, nil, err
        }
        certBytes := ssh.MarshalAuthorizedKey(sshCert)
        // Remove newline character that MarshalAuthorizedKey() adds
        certBytes = certBytes[:len(certBytes)-1]

        seckeySsh, err := ssh.MarshalPrivateKey(signer, "openpubkey cert")
        if err != nil {
                return nil, nil, err
        }
        seckeySshBytes := pem.EncodeToMemory(seckeySsh)

        return certBytes, seckeySshBytes, nil
}

func writeKeysToSSHDir(seckeySshPem []byte, certBytes []byte) error {
        homePath, err := os.UserHomeDir()
        if err != nil {
                return err
        }
        sshPath := filepath.Join(homePath, ".ssh")

        // Make ~/.ssh if folder does not exist
        err = os.MkdirAll(sshPath, os.ModePerm)
        if err != nil {
                return err
        }

        // For ssh to automatically find the key created by openpubkey when
        // connecting, we use one of the default ssh key paths. However, the file
        // might contain an existing key. We will overwrite the key if it was
        // generated by openpubkey  which we check by looking at the associated
        // comment. If the comment is equal to "openpubkey", we overwrite the file
        // with a new key.
        for _, keyFilename := range []string{"id_ecdsa", "id_ed25519"} {
                seckeyPath := filepath.Join(sshPath, keyFilename)
                pubkeyPath := seckeyPath + ".pub"

                if !fileExists(seckeyPath) {
                        // If ssh key file does not currently exist, we don't have to worry about overwriting it
                        return writeKeys(seckeyPath, pubkeyPath, seckeySshPem, certBytes)
                } else if !fileExists(pubkeyPath) {
                        continue
                } else {
                        // If the ssh key file does exist, check if it was generated by openpubkey, if it was then it is safe to overwrite
                        sshPubkey, err := os.ReadFile(pubkeyPath)
                        if err != nil {
                                log.Println("Failed to read:", pubkeyPath)
                                continue
                        }
                        _, comment, _, _, err := ssh.ParseAuthorizedKey(sshPubkey)
                        if err != nil {
                                log.Println("Failed to parse:", pubkeyPath)
                                continue
                        }

                        // If the key comment is "openpubkey" then we generated it
                        if comment == "openpubkey" {
                                return writeKeys(seckeyPath, pubkeyPath, seckeySshPem, certBytes)
                        }
                }
        }
        return fmt.Errorf("no default ssh key file free for openpubkey")
}

func writeKeys(seckeyPath string, pubkeyPath string, seckeySshPem []byte, certBytes []byte) error {
        // Write ssh secret key to filesystem
        if err := os.WriteFile(seckeyPath, seckeySshPem, 0600); err != nil {
                return err
        }

        fmt.Printf("Writing opk ssh public key to %s and corresponding secret key to %s\n", pubkeyPath, seckeyPath)

        certBytes = append(certBytes, []byte(" openpubkey")...)
        // Write ssh public key (certificate) to filesystem
        return os.WriteFile(pubkeyPath, certBytes, 0644)
}

func fileExists(fPath string) bool {
        _, err := os.Open(fPath)
        return !errors.Is(err, os.ErrNotExist)
}

func IdentityString(pkt pktoken.PKToken) (string, error) {
        idt, err := oidc.NewJwt(pkt.OpToken)
        if err != nil {
                return "", err
        }
        claims := idt.GetClaims()
        if claims.Email == "" {
                return "Sub, issuer, audience: \n" + claims.Subject + " " + claims.Issuer + " " + claims.Audience, nil
        } else {
                return "Email, sub, issuer, audience: \n" + claims.Email + " " + claims.Subject + " " + claims.Issuer + " " + claims.Audience, nil
        }
}

func PrettyIdToken(pkt pktoken.PKToken) (string, error) {

        idt, err := oidc.NewJwt(pkt.OpToken)
        if err != nil {
                return "", err
        }

        idt_json, err := json.MarshalIndent(idt.GetClaims(), "", "    ")

        if err != nil {
                return "", err
        }

        return string(idt_json[:]), 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

//go:build linux || darwin

package commands

import (
        "errors"
        "fmt"
        "os"
        "os/user"
        "path/filepath"
        "regexp"
        "strconv"
        "syscall"

        "github.com/openpubkey/opkssh/policy/files"
)

// ReadHome is used to read the home policy file for the user with
// the specified username. This is used when opkssh is called by
// AuthorizedKeysCommand as the opksshuser and needs to use sudoer
// access to read the home policy file (`/home/<username>/opk/auth_id`).
// This function is only available on Linux and Darwin because it relies on
// syscall.Stat_t to determine the owner of the file.
func ReadHome(username string) ([]byte, error) {
        if matched, _ := regexp.MatchString("^[a-z0-9_\\-.]+$", username); !matched {
                return nil, fmt.Errorf("%s is not a valid linux username", username)
        }

        userObj, err := user.Lookup(username)
        if err != nil {
                return nil, fmt.Errorf("failed to find user %s", username)
        }
        homePolicyPath := filepath.Join(userObj.HomeDir, ".opk", "auth_id")

        // Security critical: We reading this file as `sudo -u opksshuser`
        // and opksshuser has elevated permissions to read any file whose
        // path matches `/home/*/opk/auth_id`. We need to be cautious we do follow
        // a symlink as it could be to a file the user is not permitted to read.
        // This would not permit the user to read the file, but they might be able
        // to determine the existence of the file. We use O_NOFOLLOW to prevent
        // following symlinks.
        file, err := os.OpenFile(homePolicyPath, os.O_RDONLY|syscall.O_NOFOLLOW, 0)
        if err != nil {
                if errors.Is(err, syscall.ELOOP) {
                        return nil, fmt.Errorf("home policy file %s is a symlink, symlink are unsafe in this context", homePolicyPath)
                }
                return nil, fmt.Errorf("failed to open %s, %v", homePolicyPath, err)
        }
        defer file.Close()

        if fileInfo, err := file.Stat(); err != nil {
                return nil, fmt.Errorf("failed to get info on file %s", homePolicyPath)
        } else if stat, ok := fileInfo.Sys().(*syscall.Stat_t); !ok { // This syscall.Stat_t is doesn't work on Windows
                return nil, fmt.Errorf("failed to stat file %s", homePolicyPath)
        } else {
                // We want to ensure that the file is owned by the correct user and has the correct permissions.
                requiredOwnerUid := userObj.Uid
                fileOwnerUID := strconv.FormatUint(uint64(stat.Uid), 10)
                fileOwner, err := user.LookupId(fileOwnerUID)
                if err != nil {
                        return nil, fmt.Errorf("failed to find username for UID %s for file %s", fileOwnerUID, homePolicyPath)
                }
                if fileOwnerUID != userObj.Uid || fileOwner.Username != username {
                        return nil, fmt.Errorf("unsafe file permissions on %s expected file owner %s (UID %s) got %s (UID %s)",
                                homePolicyPath, username, requiredOwnerUid, fileOwner.Username, fileOwnerUID)
                }
                if fileInfo.Mode().Perm() != files.ModeHomePerms {
                        return nil, fmt.Errorf("unsafe file permissions for %s got %o expected %o", homePolicyPath, fileInfo.Mode().Perm(), files.ModeHomePerms)
                }
                fileBytes, err := os.ReadFile(homePolicyPath)
                if err != nil {
                        return nil, fmt.Errorf("failed to read %s, %v", homePolicyPath, err)
                }
                return fileBytes, 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 commands

import (
        "context"

        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/verifier"
        "github.com/openpubkey/opkssh/policy"
        "github.com/openpubkey/opkssh/sshcert"
        "golang.org/x/crypto/ssh"
)

// PolicyEnforcerFunc returns nil if the supplied PK token is permitted to login as
// username. Otherwise, an error is returned indicating the reason for rejection
type PolicyEnforcerFunc func(username string, pkt *pktoken.PKToken) error

// VerifyCmd provides functionality to verify OPK tokens contained in SSH
// certificates and authorize requests to SSH as a specific username using a
// configurable authorization system. It is designed to be used in conjunction
// with sshd's AuthorizedKeysCommand feature.
type VerifyCmd struct {
        // PktVerifier is responsible for verifying the PK token
        // contained in the SSH certificate
        PktVerifier verifier.Verifier
        // CheckPolicy determines whether the verified PK token is permitted to SSH as a
        // specific user
        CheckPolicy PolicyEnforcerFunc
}

// This function is called by the SSH server as the AuthorizedKeysCommand:
//
// The following lines are added to /etc/ssh/sshd_config:
//
//        AuthorizedKeysCommand /usr/local/bin/opkssh ver %u %k %t
//        AuthorizedKeysCommandUser opksshuser
//
// The parameters specified in the config map the parameters sent to the function below.
// We prepend "Arg" to specify which ones are arguments sent by sshd. They are:
//
//        %u The username (requested principal) - userArg
//        %k The base64-encoded public key for authentication - certB64Arg - the public key is also a certificate
//        %t The public key type - typArg - in this case a certificate being used as a public key
//
// AuthorizedKeysCommand verifies the OPK PK token contained in the base64-encoded SSH pubkey;
// the pubkey is expected to be an SSH certificate. pubkeyType is used to
// determine how to parse the pubkey as one of the SSH certificate types.
//
// This function:
// 1. Verifying the PK token with the OP (OpenID Provider)
// 2. Enforcing policy by checking if the identity is allowed to assume
// the username (principal) requested.
//
// If all steps of verification succeed, then the expected authorized_keys file
// format string is returned (i.e. the expected line to produce on standard
// output when using sshd's AuthorizedKeysCommand feature). Otherwise, a non-nil
// error is returned.
func (v *VerifyCmd) AuthorizedKeysCommand(ctx context.Context, userArg string, typArg string, certB64Arg string) (string, error) {
        // Parse the b64 pubkey and expect it to be an ssh certificate
        cert, err := sshcert.NewFromAuthorizedKey(typArg, certB64Arg)
        if err != nil {
                return "", err
        }
        if pkt, err := cert.VerifySshPktCert(ctx, v.PktVerifier); err != nil { // Verify the PKT contained in the cert
                return "", err
        } else if err := v.CheckPolicy(userArg, pkt); err != nil { // Check if username is authorized
                return "", err
        } else { // Success!
                // sshd expects the public key in the cert, not the cert itself. This
                // public key is key of the CA that signs the cert, in our setting there
                // is no CA.
                pubkeyBytes := ssh.MarshalAuthorizedKey(cert.SshCert.SignatureKey)
                return "cert-authority " + string(pubkeyBytes), nil
        }
}

// OpkPolicyEnforcerAuthFunc returns an opkssh policy.Enforcer that can be
// used in the opkssh verify command.
func OpkPolicyEnforcerFunc(username string) PolicyEnforcerFunc {
        policyEnforcer := &policy.Enforcer{
                PolicyLoader: &policy.MultiPolicyLoader{
                        HomePolicyLoader:   policy.NewHomePolicyLoader(),
                        SystemPolicyLoader: policy.NewSystemPolicyLoader(),
                        Username:           username,
                        LoadWithScript:     true, // This is needed to load policy from the user's home directory
                },
        }
        return policyEnforcer.CheckPolicy
}

// 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 main

import (
        "context"
        "errors"
        "fmt"
        "log"
        "os"
        "os/exec"
        "os/signal"
        "regexp"
        "strings"
        "syscall"

        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/opkssh/commands"
        "github.com/openpubkey/opkssh/policy"
        "github.com/openpubkey/opkssh/policy/files"
        "github.com/spf13/cobra"
)

var (
        // These can be overridden at build time using ldflags. For example:
        // go build -v -o /usr/local/bin/opkssh -ldflags "-X main.issuer=http://oidc.local:${ISSUER_PORT}/ -X main.clientID=web -X main.clientSecret=secret"
        Version           = "unversioned"
        issuer            = ""
        clientID          = ""
        clientSecret      = ""
        redirectURIs      = ""
        logFilePathServer = "/var/log/opkssh.log" // Remember if you change this, change it in the install script as well
)

func main() {
        os.Exit(run())
}

func run() int {
        rootCmd := &cobra.Command{
                SilenceUsage: true,
                Use:          "opkssh",
                Short:        "SSH with OpenPubkey",
                Version:      Version,
                Long: `SSH with OpenPubkey

This program allows users to:
  - Login and create SSH key pairs using their OpenID Connect identity
  - Add policies to auth_id policy files
  - Verify OpenPubkey SSH certificates for use with sshd's AuthorizedKeysCommand`,
                Example: `  opkssh login
  opkssh add root alice@example.com https://accounts.google.com`,
                RunE: func(cmd *cobra.Command, args []string) error {
                        return cmd.Help()
                },
        }
        rootCmd.CompletionOptions.DisableDefaultCmd = true

        addCmd := &cobra.Command{
                SilenceUsage: true,
                Use:          "add <PRINCIPAL> <EMAIL|SUB|GROUP> <ISSUER>",
                Short:        "Appends new rule to the policy file",
                Long: `Add appends a new policy entry in the auth_id policy file granting SSH access to the specified email or subscriber ID (sub) or group.

It first attempts to write to the system-wide file (/etc/opk/auth_id). If it lacks permissions to update this file it falls back to writing to the user-specific file (~/.opk/auth_id).

Arguments:
  PRINCIPAL            The target user account (requested principal).
  EMAIL|SUB|GROUP      Email address, subscriber ID or group authorized to assume this principal. If using an OIDC group, the argument needs to be in the format of oidc:groups:<groupId>.
  ISSUER               OpenID Connect provider (issuer) URL associated with the email/sub/group.
`,
                Args: cobra.ExactArgs(3),
                Example: `  opkssh add root alice@example.com https://accounts.google.com
  opkssh add alice 103030642802723203118 https://accounts.google.com
  opkssh add developer oidc:groups:developer https://accounts.google.com`,
                RunE: func(cmd *cobra.Command, args []string) error {
                        inputPrincipal := args[0]
                        inputEmail := args[1]
                        inputIssuer := args[2]

                        // Convenience aliases to save user time (who is going to remember the hideous Azure issuer string)
                        switch inputIssuer {
                        case "google":
                                inputIssuer = "https://accounts.google.com"
                        case "azure", "microsoft":
                                inputIssuer = "https://login.microsoftonline.com/9188040d-6c67-4c5b-b112-36a304b66dad/v2.0"
                        case "gitlab":
                                inputIssuer = "https://gitlab.com"
                        }

                        add := commands.AddCmd{
                                HomePolicyLoader:   policy.NewHomePolicyLoader(),
                                SystemPolicyLoader: policy.NewSystemPolicyLoader(),
                                Username:           inputPrincipal,
                        }
                        policyFilePath, err := add.Run(inputPrincipal, inputEmail, inputIssuer)
                        if err != nil {
                                fmt.Fprintf(os.Stderr, "Failed to add to policy: %v\n", err)
                                return err
                        }
                        fmt.Fprintf(os.Stdout, "Successfully added new policy to %s\n", policyFilePath)
                        return nil
                },
        }
        rootCmd.AddCommand(addCmd)

        var autoRefresh bool
        var logDir string
        var providerArg string
        var disableBrowserOpenArg bool
        var printIdTokenArg bool
        loginCmd := &cobra.Command{
                SilenceUsage: true,
                Use:          "login",
                Short:        "Authenticate with an OpenID Provider to generate an SSH key for opkssh",
                Long: `Login creates opkssh SSH keys

Login generates a key pair, then opens a browser to authenticate the user with the OpenID Provider. Upon successful authentication, opkssh creates an SSH public key (~/.ssh/id_ecdsa) containing the user's PK token. By default, this SSH key expires after 24 hours, after which the user must run "opkssh login" again to generate a new key.

Users can then SSH into servers configured to use opkssh as the AuthorizedKeysCommand. The server verifies the PK token and grants access if the token is valid and the user is authorized per the auth_id policy.
`,
                Example: `  opkssh login
  opkssh login --provider=<issuer>,<client_id>`,
                RunE: func(cmd *cobra.Command, args []string) error {

                        ctx, cancel := context.WithCancel(context.Background())
                        defer cancel()
                        sigs := make(chan os.Signal, 1)
                        signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
                        go func() {
                                <-sigs
                                cancel()
                        }()

                        // If LDFlags issuer is set, build providerFromLdFlags
                        var providerFromLdFlags providers.OpenIdProvider
                        if issuer != "" {
                                opts := providers.GetDefaultGoogleOpOptions()
                                opts.Issuer = issuer
                                opts.ClientID = clientID
                                opts.ClientSecret = clientSecret
                                opts.RedirectURIs = strings.Split(redirectURIs, ",")
                                providerFromLdFlags = providers.NewGoogleOpWithOptions(opts)
                        }

                        login := commands.NewLogin(autoRefresh, logDir, disableBrowserOpenArg, printIdTokenArg, providerArg, providerFromLdFlags)
                        if err := login.Run(ctx); err != nil {
                                log.Println("Error executing login command:", err)
                                return err
                        }
                        return nil
                },
        }

        // Define flags for login.
        loginCmd.Flags().BoolVar(&autoRefresh, "auto-refresh", false, "Automatically refresh PK token after login")
        loginCmd.Flags().StringVar(&logDir, "log-dir", "", "Directory to write output logs")
        loginCmd.Flags().BoolVar(&disableBrowserOpenArg, "disable-browser-open", false, "Set this flag to disable opening the browser. Useful for choosing the browser you want to use.")
        loginCmd.Flags().BoolVar(&printIdTokenArg, "print-id-token", false, "Set this flag to print out the contents of the id_token. Useful for inspecting claims.")
        loginCmd.Flags().StringVar(&providerArg, "provider", "", "OpenID Provider specification in the format: <issuer>,<client_id> or <issuer>,<client_id>,<client_secret>")
        rootCmd.AddCommand(loginCmd)

        readhomeCmd := &cobra.Command{
                SilenceUsage: true,
                Use:          "readhome <PRINCIPAL>",
                Short:        "Read the principal's home policy file",
                Long: `Read the principal's policy file (/home/<PRINCIPAL>/.opk/auth_id).

You should not call this command directly. It is called by the opkssh verify command as part of the AuthorizedKeysCommand process to read the user's policy  (principals) home file (~/.opk/auth_id) with sudoer permissions. This allows us to use an unprivileged user as the AuthorizedKeysCommand user.
`,
                Args:    cobra.ExactArgs(1),
                Example: `  opkssh readhome alice`,
                RunE: func(cmd *cobra.Command, args []string) error {
                        userArg := os.Args[2]
                        if fileBytes, err := commands.ReadHome(userArg); err != nil {
                                fmt.Fprintf(os.Stderr, "Failed to read user's home policy file: %v\n", err)
                                return err
                        } else {
                                fmt.Fprint(os.Stdout, string(fileBytes))
                                return nil
                        }
                },
        }
        rootCmd.AddCommand(readhomeCmd)

        verifyCmd := &cobra.Command{
                SilenceUsage: true,
                Use:          "verify <PRINCIPAL> <CERT> <KEY_TYPE>",
                Short:        "Verify an SSH key (used by sshd AuthorizedKeysCommand)",
                Long: `Verify extracts a PK token from a base64-encoded SSH certificate and verifies it against policy. It expects an allowed provider file at /etc/opk/providers and a user policy file at either /etc/opk/auth_id or ~/.opk/auth_id.

This command is intended to be called by sshd as an AuthorizedKeysCommand:
  https://man.openbsd.org/sshd_config#AuthorizedKeysCommand

During installation, opkssh typically adds these lines to /etc/ssh/sshd_config:
  AuthorizedKeysCommand /usr/local/bin/opkssh verify %%u %%k %%t
  AuthorizedKeysCommandUser opksshuser

Where the tokens in /etc/ssh/sshd_config are defined as:
  %%u   Target username (requested principal)
  %%k   Base64-encoded SSH public key (SSH certificate) provided for authentication
  %%t   Public key type (SSH certificate format, e.g., ecdsa-sha2-nistp256-cert-v01@openssh.com)

Verification checks performed:
  1. Ensures the PK token is properly formed, signed, and issued by the specified OpenID Provider (OP).
  2. Confirms the PK token's issue (iss) and client ID (audience) are listed in the allowed provider file (/etc/opk/providers) and the token is not expired.
  3. Validates the identity (email or sub) in the PK token against user policies (/etc/opk/auth_id or ~/.opk/auth_id) to ensure it can assume the requested username (principal).

If all checks pass, Verify authorizes the SSH connection.

Arguments:
  PRINCIPAL    Target username.
  CERT         Base64-encoded SSH certificate.
  KEY_TYPE     SSH certificate key type (e.g., ecdsa-sha2-nistp256-cert-v01@openssh.com)`,
                Args:    cobra.ExactArgs(3),
                Example: `  opkssh verify root <base64-encoded-cert> ecdsa-sha2-nistp256-cert-v01@openssh.com`,
                RunE: func(cmd *cobra.Command, args []string) error {
                        ctx := context.Background()

                        // Setup logger
                        logFile, err := os.OpenFile(logFilePathServer, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0660) // Owner and group can read/write
                        if err != nil {
                                fmt.Fprintf(os.Stderr, "Error opening log file: %v\n", err)
                                // It could be very difficult to figure out what is going on if the log file was deleted. Hopefully this message saves someone an hour of debugging.
                                fmt.Fprintf(os.Stderr, "Check if log exists at %v, if it does not create it with permissions: chown root:opksshuser %v; chmod 660 %v\n", logFilePathServer, logFilePathServer, logFilePathServer)
                        } else {
                                defer logFile.Close()
                                log.SetOutput(logFile)
                        }

                        // Logs if using an unsupported OpenSSH version
                        checkOpenSSHVersion()

                        // The "AuthorizedKeysCommand" func is designed to be used by sshd and specified as an AuthorizedKeysCommand
                        // ref: https://man.openbsd.org/sshd_config#AuthorizedKeysCommand
                        log.Println(strings.Join(os.Args, " "))

                        userArg := args[0]
                        certB64Arg := args[1]
                        typArg := args[2]

                        providerPolicyPath := "/etc/opk/providers"
                        providerPolicy, err := policy.NewProviderFileLoader().LoadProviderPolicy(providerPolicyPath)

                        if err != nil {
                                log.Println("Failed to open /etc/opk/providers:", err)
                                return err
                        }

                        printConfigProblems()
                        log.Println("Providers loaded: ", providerPolicy.ToString())

                        pktVerifier, err := providerPolicy.CreateVerifier()
                        if err != nil {
                                log.Println("Failed to create pk token verifier (likely bad configuration):", err)
                                return err
                        }

                        v := commands.VerifyCmd{
                                PktVerifier: *pktVerifier,
                                CheckPolicy: commands.OpkPolicyEnforcerFunc(userArg),
                        }
                        if authKey, err := v.AuthorizedKeysCommand(ctx, userArg, typArg, certB64Arg); err != nil {
                                log.Println("failed to verify:", err)
                                return err
                        } else {
                                log.Println("successfully verified")
                                // sshd is awaiting a specific line, which we print here. Printing anything else before or after will break our solution
                                fmt.Println(authKey)
                                return nil
                        }
                },
        }
        rootCmd.AddCommand(verifyCmd)

        err := rootCmd.Execute()
        if err != nil {
                return 1
        }
        return 0
}

func printConfigProblems() {
        problems := files.ConfigProblems().GetProblems()
        if len(problems) > 0 {
                log.Println("Warning: Encountered the following configuration problems:")
                for _, problem := range problems {
                        log.Println(problem.String())
                }
        }
}

// OpenSSH used to impose a 4096-octet limit on the string buffers available to
// the percent_expand function. In October 2019 as part of the 8.1 release,
// that limit was removed. If you exceeded this amount it would fail with
// fatal: percent_expand: string too long
// The following two functions check whether the OpenSSH version on the
// system running the verifier is greater than or equal to 8.1;
// if not then prints a warning
func checkOpenSSHVersion() {

        // Redhat/centos does not recognize `sshd -V` but does recognize `ssh -V`
        // Ubuntu recognizes both
        cmd := exec.Command("ssh", "-V")
        output, err := cmd.CombinedOutput()
        if err != nil {
                log.Println("Warning: Error executing ssh -V:", err)
                return
        }

        if ok, _ := isOpenSSHVersion8Dot1OrGreater(string(output)); !ok {
                log.Println("Warning: OpenPubkey SSH requires OpenSSH v. 8.1 or greater")
        }
}

func isOpenSSHVersion8Dot1OrGreater(opensshVersion string) (bool, error) {
        // To handle versions like 9.9p1; we only need the initial numeric part for the comparison
        re, err := regexp.Compile(`^(\d+(?:\.\d+)*).*`)
        if err != nil {
                fmt.Println("Error compiling regex:", err)
                return false, err
        }

        opensshVersion = strings.TrimPrefix(
                strings.Split(opensshVersion, ", ")[0],
                "OpenSSH_",
        )

        matches := re.FindStringSubmatch(opensshVersion)

        if len(matches) <= 0 {
                fmt.Println("Invalid OpenSSH version")
                return false, errors.New("invalid OpenSSH version")
        }

        version := matches[1]

        if version >= "8.1" {
                return true, nil
        }

        return false, 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 policy

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

        "github.com/openpubkey/openpubkey/pktoken"
        "golang.org/x/exp/slices"
)

// Enforcer evaluates opkssh policy to determine if the desired principal is
// permitted
type Enforcer struct {
        PolicyLoader Loader
}

// type for Identity Token checkedClaims
type checkedClaims struct {
        Email  string   `json:"email"`
        Sub    string   `json:"sub"`
        Groups []string `json:"groups"`
}

// Validates that the server defined identity attribute matches the
// respective claim from the identity token
func validateClaim(claims *checkedClaims, user *User) bool {
        if strings.HasPrefix(user.IdentityAttribute, "oidc:groups") {
                oidcGroupSections := strings.Split(user.IdentityAttribute, ":")

                return slices.Contains(claims.Groups, oidcGroupSections[len(oidcGroupSections)-1])
        }

        // email should be a case-insensitive check
        // sub should be a case-sensitive check
        return strings.EqualFold(claims.Email, user.IdentityAttribute) || string(claims.Sub) == user.IdentityAttribute
}

// CheckPolicy loads the opkssh policy and checks to see if there is a policy
// permitting access to principalDesired for the user identified by the PKT's
// email claim. Returns nil if access is granted. Otherwise, an error is
// returned.
//
// It is security critical to verify the pkt first before calling this function.
// This is because if this function is called first, a timing channel exists which
// allows an attacker check what identities and principals are allowed by the policy.
func (p *Enforcer) CheckPolicy(principalDesired string, pkt *pktoken.PKToken) error {
        policy, source, err := p.PolicyLoader.Load()
        if err != nil {
                return fmt.Errorf("error loading policy: %w", err)
        }

        sourceStr := source.Source()
        if sourceStr == "" {
                sourceStr = "<policy source unknown>"
        }

        var claims checkedClaims

        if err := json.Unmarshal(pkt.Payload, &claims); err != nil {
                return fmt.Errorf("error unmarshalling pk token payload: %w", err)
        }
        issuer, err := pkt.Issuer()
        if err != nil {
                return fmt.Errorf("error getting issuer from pk token: %w", err)
        }
        for _, user := range policy.Users {
                // check each entry to see if the user in the checkedClaims is included
                if validateClaim(&claims, &user) {
                        if issuer != user.Issuer {
                                continue
                        }
                        // if they are, then check if the desired principal is allowed
                        if slices.Contains(user.Principals, principalDesired) {
                                // access granted
                                return nil
                        }
                }
        }

        return fmt.Errorf("no policy to allow %s with (issuer=%s) to assume %s, check policy config at %s", claims.Email, issuer, principalDesired, sourceStr)
}

// 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 files

import (
        "fmt"
        "strings"
        "sync"
)

type ConfigProblem struct {
        Filepath            string
        OffendingLine       string
        OffendingLineNumber int
        ErrorMessage        string
        Source              string
}

func (e ConfigProblem) String() string {
        return "encountered error: " + e.ErrorMessage + ", reading " + e.OffendingLine + " in " + e.Filepath + " at line " + fmt.Sprint(e.OffendingLineNumber)
}

type ConfigLog struct {
        log      []ConfigProblem
        logMutex sync.Mutex
}

func (c *ConfigLog) RecordProblem(entry ConfigProblem) {
        c.logMutex.Lock()
        defer c.logMutex.Unlock()
        c.log = append(c.log, entry)
}

func (c *ConfigLog) GetProblems() []ConfigProblem {
        c.logMutex.Lock()
        defer c.logMutex.Unlock()
        logCopy := make([]ConfigProblem, len(c.log))
        copy(logCopy, c.log)
        return logCopy
}

func (c *ConfigLog) NoProblems() bool {
        c.logMutex.Lock()
        defer c.logMutex.Unlock()
        return len(c.log) == 0
}

func (c *ConfigLog) String() string {
        // No mutex needed since GetLogs handles the mutex
        logs := c.GetProblems()
        logsStrings := []string{}
        for _, log := range logs {
                logsStrings = append(logsStrings, log.String())
        }
        return strings.Join(logsStrings, "\n")
}

func (c *ConfigLog) Clear() {
        c.logMutex.Lock()
        defer c.logMutex.Unlock()
        c.log = []ConfigProblem{}
}

var (
        singleton *ConfigLog
        once      sync.Once
)

func ConfigProblems() *ConfigLog {
        once.Do(func() {
                singleton = &ConfigLog{
                        log:      []ConfigProblem{},
                        logMutex: sync.Mutex{},
                }
        })
        return singleton
}

// 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 files

import (
        "fmt"
        "io/fs"
        "path/filepath"

        "github.com/spf13/afero"
)

// UserPolicyLoader contains methods to read/write the opkssh policy file from/to an
// arbitrary filesystem. All methods that read policy from the filesystem fail
// and return an error immediately if the permission bits are invalid.
type FileLoader struct {
        Fs           afero.Fs
        RequiredPerm fs.FileMode
}

// CreateIfDoesNotExist creates a file at the given path if it does not exist.
func (l FileLoader) CreateIfDoesNotExist(path string) error {
        exists, err := afero.Exists(l.Fs, path)
        if err != nil {
                return err
        }
        if !exists {
                dirPath := filepath.Dir(path)
                if err := l.Fs.MkdirAll(dirPath, 0750); err != nil {
                        return fmt.Errorf("failed to create directory: %w", err)
                }
                file, err := l.Fs.Create(path)
                if err != nil {
                        return fmt.Errorf("failed to create file: %w", err)
                }
                file.Close()
                if err := l.Fs.Chmod(path, l.RequiredPerm); err != nil {
                        return fmt.Errorf("failed to set file permissions: %w", err)
                }
        }
        return nil
}

// LoadFileAtPath validates that the file at path exists, can be read
// by the current process, and has the correct permission bits set. Parses the
// contents and returns the bytes if file permissions are valid and
// reading is successful; otherwise returns an error.
func (l *FileLoader) LoadFileAtPath(path string) ([]byte, error) {
        // Check if file exists and we can access it
        if _, err := l.Fs.Stat(path); err != nil {
                return nil, fmt.Errorf("failed to describe the file at path: %w", err)
        }

        // Validate that file has correct permission bits set
        if err := NewPermsChecker(l.Fs).CheckPerm(path, l.RequiredPerm, "", ""); err != nil {
                return nil, fmt.Errorf("policy file has insecure permissions: %w", err)
        }

        // Read file contents
        afs := &afero.Afero{Fs: l.Fs}
        content, err := afs.ReadFile(path)
        if err != nil {
                return nil, err
        }
        return content, nil
}

// Dump writes the bytes in fileBytes to the filepath
func (l *FileLoader) Dump(fileBytes []byte, path string) error {
        // Write to disk
        if err := afero.WriteFile(l.Fs, path, fileBytes, l.RequiredPerm); err != nil {
                return err
        }
        return 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 files

import (
        "fmt"
        "io/fs"
        "log"
        "os/exec"
        "strings"

        "github.com/spf13/afero"
)

// ModeSystemPerms is the expected permission bits that should be set for opkssh
// system policy files (`/etc/opk/auth_id`, `/etc/opk/providers`). This mode means
// that only the owner of the file can write/read to the file, but the group which
// should be opksshuser can read the file.
const ModeSystemPerms = fs.FileMode(0640)

// ModeHomePerms is the expected permission bits that should be set for opkssh
// user home policy files `~/.opk/auth_id`.
const ModeHomePerms = fs.FileMode(0600)

// PermsChecker contains methods to check the ownership, group
// and file permissions of a file on a Unix-like system.
type PermsChecker struct {
        Fs        afero.Fs
        cmdRunner func(string, ...string) ([]byte, error)
}

func NewPermsChecker(fs afero.Fs) *PermsChecker {
        return &PermsChecker{Fs: fs, cmdRunner: execCmd}
}

// CheckPerm checks the file at the given path if it has the desired permissions.
// If the requiredOwner or requiredGroup are not empty then the function will also
// that the owner and group of the file match the requiredOwner and requiredGroup
// specified and fail if they do not.
func (u *PermsChecker) CheckPerm(path string, requirePerm fs.FileMode, requiredOwner string, requiredGroup string) error {
        fileInfo, err := u.Fs.Stat(path)
        if err != nil {
                return fmt.Errorf("failed to describe the file at path: %w", err)
        }
        mode := fileInfo.Mode()

        // if the requiredOwner or requiredGroup are specified then run stat and check if they match
        if requiredOwner != "" || requiredGroup != "" {
                log.Println("Running, command: ", "stat", "-c", "%U %G", path)
                statOutput, err := u.cmdRunner("stat", "-c", "%U %G", path)
                log.Println("Got output:", string(statOutput))
                if err != nil {
                        return fmt.Errorf("failed to run stat: %w", err)
                }

                statOutputSplit := strings.Split(strings.TrimSpace(string(statOutput)), " ")
                statOwner := statOutputSplit[0]
                statGroup := statOutputSplit[1]
                if len(statOutputSplit) != 2 {
                        return fmt.Errorf("expected stat command to return 2 values got %d", len(statOutputSplit))
                }

                if requiredOwner != "" {
                        if requiredOwner != statOwner {
                                return fmt.Errorf("expected owner (%s), got (%s)", requiredOwner, statOwner)
                        }
                }
                if requiredGroup != "" {
                        if requiredGroup != statGroup {
                                return fmt.Errorf("expected group (%s), got (%s)", requiredGroup, statGroup)
                        }
                }
        }

        if mode.Perm() != requirePerm {
                return fmt.Errorf("expected permissions (%o), got (%o)", requirePerm.Perm(), mode.Perm())
        }

        return nil
}

func execCmd(name string, arg ...string) ([]byte, error) {
        cmd := exec.Command(name, arg...)
        return cmd.CombinedOutput()
}

// 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 files

import "strings"

type Table struct {
        rows [][]string
}

func NewTable(content []byte) *Table {
        table := [][]string{}
        rows := strings.Split(string(content), "\n")
        for _, row := range rows {
                row := CleanRow(row)
                if row == "" {
                        continue
                }
                columns := strings.Fields(row)
                table = append(table, columns)
        }
        return &Table{rows: table}
}

func CleanRow(row string) string {
        // Remove comments
        rowFixed := strings.Split(row, "#")[0]
        // Skip empty rows
        rowFixed = strings.TrimSpace(rowFixed)
        return rowFixed
}

func (t *Table) AddRow(row ...string) {
        t.rows = append(t.rows, row)
}

func (t Table) ToString() string {
        var sb strings.Builder
        for _, row := range t.rows {
                sb.WriteString(strings.Join(row, " ") + "\n")
        }
        return sb.String()
}

func (t Table) ToBytes() []byte {
        return []byte(t.ToString())
}

func (t Table) GetRows() [][]string {
        return t.rows
}

// 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 policy

import (
        "errors"
        "fmt"
        "log"
        "os/exec"
        "strings"
)

var _ Loader = &MultiPolicyLoader{}

// FileSource implements policy.Source by returning a string that is expected to
// be a filepath
type FileSource string

func (s FileSource) Source() string {
        return string(s)
}

// MultiPolicyLoader implements policy.Loader by reading both the system default
// policy (root policy) and user policy (~/.opk/auth_id where ~ maps to
// Username's home directory)
type MultiPolicyLoader struct {
        HomePolicyLoader   *HomePolicyLoader
        SystemPolicyLoader *SystemPolicyLoader
        LoadWithScript     bool
        Username           string
}

func (l *MultiPolicyLoader) Load() (*Policy, Source, error) {
        policy := new(Policy)

        // Try to load the root policy
        rootPolicy, _, rootPolicyErr := l.SystemPolicyLoader.LoadSystemPolicy()
        if rootPolicyErr != nil {
                log.Println("warning: failed to load system default policy:", rootPolicyErr)
        }
        // Try to load the user policy
        userPolicy, userPolicyFilePath, userPolicyErr := l.HomePolicyLoader.LoadHomePolicy(l.Username, true, ReadWithSudoScript)
        if userPolicyErr != nil {
                log.Println("warning: failed to load user policy:", userPolicyErr)
        }
        // Log warning if no error loading, but userPolicy is empty meaning that
        // there are no valid entries
        if userPolicyErr == nil && len(userPolicy.Users) == 0 {
                log.Printf("warning: user policy %s has no valid user entries; an entry is considered valid if it gives %s access.", userPolicyFilePath, l.Username)
        }

        // Failed to read both policies. Return multi-error
        if rootPolicy == nil && userPolicy == nil {
                return nil, EmptySource{}, errors.Join(rootPolicyErr, userPolicyErr)
        }

        // TODO-Yuval: Optimize by merging duplicate entries instead of blindly
        // appending
        readPaths := []string{}
        if rootPolicy != nil {
                policy.Users = append(policy.Users, rootPolicy.Users...)
                readPaths = append(readPaths, SystemDefaultPolicyPath)
        }
        if userPolicy != nil {
                policy.Users = append(policy.Users, userPolicy.Users...)
                readPaths = append(readPaths, userPolicyFilePath)
        }

        return policy, FileSource(strings.Join(readPaths, ", ")), nil
}

// ReadWithSudoScript specifies additional way of loading the policy in the
// user's home directory (`~/.opk/auth_id`). This is needed when the
// AuthorizedKeysCommand user does not have privileges to transverse the user's
// home directory. Instead we call run a command which uses special
// sudoers permissions to read the policy file.
//
// Doing this is more secure than simply giving opkssh sudoer access because
// if there was an RCE in opkssh could be triggered an SSH request via
// AuthorizedKeysCommand, the new opkssh process we use to perform the read
// would not be compromised. Thus, the compromised opkssh process could not assume
// full root privileges.
func ReadWithSudoScript(h *HomePolicyLoader, username string) ([]byte, error) {
        // opkssh readhome ensures the file is not a symlink and has the permissions/ownership.
        cmd := exec.Command("sudo", "-n", "/usr/local/bin/opkssh", "readhome", username)
        homePolicyFileBytes, err := cmd.CombinedOutput()
        if err != nil {
                return nil, fmt.Errorf("error reading %s home policy using command %v got output %v and err %v", username, cmd, string(homePolicyFileBytes), err)
        }
        return homePolicyFileBytes, 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 policy

import (
        "fmt"
        "log"
        "strings"

        "github.com/openpubkey/opkssh/policy/files"
)

// User is an opkssh policy user entry
type User struct {
        // IdentityAttribute is a string that is either structured or unstructured.
        // Structured: <IdentityProtocolMatching>:<Attribute>:<Value>
        // E.g. `oidc:groups:ssh-users`
        // Using the structured identifier allows the capability of constructing
        // complex user matchers.
        //
        // Unstructured:
        // This is older version that only works with OIDC Identity Tokens, with
        // the claim being `email` or `sub`. The expected value is to be the user's
        // email or the user's subscriber ID. The expected value used when comparing
        // against an id_token's email claim Subscriber ID is a unique identifier
        // for the user at the OpenID Provider
        IdentityAttribute string
        // Principals is a list of allowed principals
        Principals []string
        // Sub        string
        Issuer string
}

// Policy represents an opkssh policy
type Policy struct {
        // Users is a list of all user entries in the policy
        Users []User
}

// FromTable decodes whitespace delimited input into policy.Policy
func FromTable(input []byte, path string) *Policy {
        table := files.NewTable(input)
        policy := &Policy{}
        for i, row := range table.GetRows() {
                // Error should not break everyone's ability to login, skip those rows
                if len(row) != 3 {
                        configProblem := files.ConfigProblem{
                                Filepath:            path,
                                OffendingLine:       strings.Join(row, " "),
                                OffendingLineNumber: i,
                                ErrorMessage:        fmt.Sprintf("wrong number of arguments (expected=3, got=%d)", len(row)),
                                Source:              "user policy file",
                        }
                        files.ConfigProblems().RecordProblem(configProblem)
                        continue
                }
                user := User{
                        Principals:        []string{row[0]},
                        IdentityAttribute: row[1],
                        Issuer:            row[2],
                }
                policy.Users = append(policy.Users, user)
        }
        return policy
}

// AddAllowedPrincipal adds a new allowed principal to the user whose email is
// equal to userEmail. If no user can be found with the email userEmail, then a
// new user entry is added with an initial allowed principals list containing
// principal. No changes are made if the principal is already allowed for this
// user.
func (p *Policy) AddAllowedPrincipal(principal string, userEmail string, issuer string) {
        userExists := false
        if len(p.Users) != 0 {
                // search to see if the current user already has an entry in the policy
                // file
                for i := range p.Users {
                        user := &p.Users[i]
                        if user.IdentityAttribute == userEmail && user.Issuer == issuer {
                                principalExists := false
                                for _, p := range user.Principals {
                                        // if the principal already exists for this user, then skip
                                        if p == principal {
                                                log.Printf("User with email %s already has access under the principal %s, skipping...\n", userEmail, principal)
                                                principalExists = true
                                        }
                                }

                                if !principalExists {
                                        user.Principals = append(user.Principals, principal)
                                        user.Issuer = issuer
                                        log.Printf("Successfully added user with email %s with principal %s to the policy file\n", userEmail, principal)
                                }
                                userExists = true
                        }
                }
        }

        // if the policy is empty or if no user found with userEmail, then create a
        // new entry
        if len(p.Users) == 0 || !userExists {
                newUser := User{
                        IdentityAttribute: userEmail,
                        Principals:        []string{principal},
                        Issuer:            issuer,
                }
                // add the new user to the list of users in the policy
                p.Users = append(p.Users, newUser)
        }
}

// ToTable encodes the policy into a whitespace delimited table
func (p *Policy) ToTable() ([]byte, error) {
        table := files.Table{}
        for _, user := range p.Users {
                for _, principal := range user.Principals {
                        table.AddRow(principal, user.IdentityAttribute, user.Issuer)
                }
        }
        return table.ToBytes(), nil
}

// Source declares the minimal interface to describe the source of a fetched
// opkssh policy (i.e. where the policy is retrieved from)
type Source interface {
        // Source returns a string describing the source of an opkssh policy. The
        // returned value is empty if there is no information about its source
        Source() string
}

var _ Source = &EmptySource{}

// EmptySource implements policy.Source and returns an empty string as the
// source
type EmptySource struct{}

func (EmptySource) Source() string { return "" }

// Loader declares the minimal interface to retrieve an opkssh policy from an
// arbitrary source
type Loader interface {
        // Load fetches an opkssh policy and returns information describing its
        // source. If an error occurs, all return values are nil except the error
        // value
        Load() (*Policy, Source, error)
}

// 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 policy

import (
        "fmt"
        "os/user"
        "path"
        "path/filepath"

        "github.com/openpubkey/opkssh/policy/files"
        "github.com/spf13/afero"
        "golang.org/x/exp/slices"
)

// SystemDefaultPolicyPath is the default filepath where opkssh policy is
// defined
var SystemDefaultPolicyPath = filepath.FromSlash("/etc/opk/auth_id")

// UserLookup defines the minimal interface to lookup users on the current
// system
type UserLookup interface {
        Lookup(username string) (*user.User, error)
}

// OsUserLookup implements the UserLookup interface by invoking the os/user
// library
type OsUserLookup struct{}

func NewOsUserLookup() UserLookup {
        return &OsUserLookup{}
}
func (OsUserLookup) Lookup(username string) (*user.User, error) { return user.Lookup(username) }

// PolicyLoader contains methods to read/write the opkssh policy file from/to an
// arbitrary filesystem. All methods that read policy from the filesystem fail
// and return an error immediately if the permission bits are invalid.
type PolicyLoader struct {
        FileLoader files.FileLoader
        UserLookup UserLookup
}

func (l PolicyLoader) CreateIfDoesNotExist(path string) error {
        return l.FileLoader.CreateIfDoesNotExist(path)
}

// LoadPolicyAtPath validates that the policy file at path exists, can be read
// by the current process, and has the correct permission bits set. Parses the
// contents and returns a policy.Policy if file permissions are valid and
// reading is successful; otherwise returns an error.
func (l *PolicyLoader) LoadPolicyAtPath(path string) (*Policy, error) {
        content, err := l.FileLoader.LoadFileAtPath(path)
        if err != nil {
                return nil, err
        }

        policy := FromTable(content, path)
        return policy, nil
}

// Dump encodes the policy into file and writes the contents to the filepath
// path
func (l *PolicyLoader) Dump(policy *Policy, path string) error {
        fileBytes, err := policy.ToTable()
        if err != nil {
                return err
        }

        // Write to disk
        if err := l.FileLoader.Dump(fileBytes, path); err != nil {
                return fmt.Errorf("failed to write to policy file %s: %w", path, err)
        }

        return nil
}

// NewSystemPolicyLoader returns an opkssh policy loader that uses the os library to
// read/write system policy from/to the filesystem.
func NewSystemPolicyLoader() *SystemPolicyLoader {
        return &SystemPolicyLoader{
                PolicyLoader: &PolicyLoader{
                        FileLoader: files.FileLoader{
                                Fs:           afero.NewOsFs(),
                                RequiredPerm: files.ModeSystemPerms,
                        },
                        UserLookup: NewOsUserLookup(),
                },
        }
}

// SystemPolicyLoader contains methods to read/write the system wide  opkssh policy file
// from/to a filesystem. All methods that read policy from the filesystem fail
// and return an error immediately if the permission bits are invalid.
type SystemPolicyLoader struct {
        *PolicyLoader
}

// LoadSystemPolicy reads the opkssh policy at SystemDefaultPolicyPath.
// An error is returned if the file cannot be read or if the permissions bits
// are not correct.
func (s *SystemPolicyLoader) LoadSystemPolicy() (*Policy, Source, error) {
        policy, err := s.LoadPolicyAtPath(SystemDefaultPolicyPath)
        if err != nil {
                return nil, EmptySource{}, fmt.Errorf("failed to read system default policy file %s: %w", SystemDefaultPolicyPath, err)
        }
        return policy, FileSource(SystemDefaultPolicyPath), nil
}

type OptionalLoader func(h *HomePolicyLoader, username string) ([]byte, error)

// HomePolicyLoader contains methods to read/write the opkssh policy file stored in
// `~/.opk/ssh` from/to a filesystem. All methods that read policy from the filesystem fail
// and return an error immediately if the permission bits are invalid.
type HomePolicyLoader struct {
        *PolicyLoader
}

// NewHomePolicyLoader returns an opkssh policy loader that uses the os library to
// read/write policy from/to the user's home directory, e.g. `~/.opk/auth_id`,
func NewHomePolicyLoader() *HomePolicyLoader {
        return &HomePolicyLoader{
                PolicyLoader: &PolicyLoader{
                        FileLoader: files.FileLoader{
                                Fs:           afero.NewOsFs(),
                                RequiredPerm: files.ModeHomePerms,
                        },
                        UserLookup: NewOsUserLookup(),
                },
        }
}

// LoadHomePolicy reads the user's opkssh policy at ~/.opk/auth_id (where ~
// maps to username's home directory) and returns the filepath read. An error is
// returned if the file cannot be read, if the permission bits are not correct,
// or if there is no user with username or has no home directory.
//
// If skipInvalidEntries is true, then invalid user entries are skipped and not
// included in the returned policy. A user policy's entry is considered valid if
// it gives username access. The returned policy is stripped of invalid entries.
// To specify an alternative Loader that will be used if we don't have sufficient
// permissions to read the policy file in the user's home directory, pass the
// alternative loader as the last argument.
func (h *HomePolicyLoader) LoadHomePolicy(username string, skipInvalidEntries bool, optLoader ...OptionalLoader) (*Policy, string, error) {
        policyFilePath, err := h.UserPolicyPath(username)
        if err != nil {
                return nil, "", fmt.Errorf("error getting user policy path for user %s: %w", username, err)
        }

        policyBytes, userPolicyErr := h.FileLoader.LoadFileAtPath(policyFilePath)
        if userPolicyErr != nil {
                if len(optLoader) == 1 {
                        // Try to read using the optional loader
                        policyBytes, err = optLoader[0](h, username)
                        if err != nil {
                                return nil, "", fmt.Errorf("failed to read user policy file %s: %w", policyFilePath, err)
                        }
                } else if len(optLoader) > 1 {
                        return nil, "", fmt.Errorf("only one optional loaders allowed, got %d", len(optLoader))
                } else {
                        return nil, "", fmt.Errorf("failed to read user policy file %s: %w", policyFilePath, userPolicyErr)
                }
        }
        policy := FromTable(policyBytes, policyFilePath)

        if skipInvalidEntries {
                // Build valid user policy. Ignore user entries that give access to a
                // principal not equal to the username where the policy file was read
                // from.
                validUserPolicy := new(Policy)
                for _, user := range policy.Users {
                        if slices.Contains(user.Principals, username) {
                                // Build clean entry that only gives access to username
                                validUserPolicy.Users = append(validUserPolicy.Users, User{
                                        IdentityAttribute: user.IdentityAttribute,
                                        Principals:        []string{username},
                                        Issuer:            user.Issuer,
                                })
                        }
                }
                return validUserPolicy, policyFilePath, nil
        } else {
                // Just return what we read
                return policy, policyFilePath, nil
        }
}

// UserPolicyPath returns the path to the user's opkssh policy file at
// ~/.opk/auth_id.
func (h *HomePolicyLoader) UserPolicyPath(username string) (string, error) {
        user, err := h.UserLookup.Lookup(username)
        if err != nil {
                return "", fmt.Errorf("failed to lookup username %s: %w", username, err)
        }
        userHomeDirectory := user.HomeDir
        if userHomeDirectory == "" {
                return "", fmt.Errorf("user %s does not have a home directory", username)
        }

        policyFilePath := path.Join(userHomeDirectory, ".opk", "auth_id")
        return policyFilePath, 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 policy

import (
        "fmt"
        "strings"

        "github.com/openpubkey/openpubkey/providers"
        "github.com/openpubkey/openpubkey/verifier"
        "github.com/openpubkey/opkssh/policy/files"
        "github.com/spf13/afero"
)

type ProvidersRow struct {
        Issuer           string
        ClientID         string
        ExpirationPolicy string
}

func (p ProvidersRow) GetExpirationPolicy() (verifier.ExpirationPolicy, error) {
        switch p.ExpirationPolicy {
        case "24h":
                return verifier.ExpirationPolicies.MAX_AGE_24HOURS, nil
        case "48h":
                return verifier.ExpirationPolicies.MAX_AGE_48HOURS, nil
        case "1week":
                return verifier.ExpirationPolicies.MAX_AGE_1WEEK, nil
        case "oidc":
                return verifier.ExpirationPolicies.OIDC, nil
        case "oidc_refreshed":
                return verifier.ExpirationPolicies.OIDC_REFRESHED, nil
        case "never":
                return verifier.ExpirationPolicies.NEVER_EXPIRE, nil
        default:
                return verifier.ExpirationPolicy{}, fmt.Errorf("invalid expiration policy: %s", p.ExpirationPolicy)
        }
}

func (p ProvidersRow) ToString() string {
        return p.Issuer + " " + p.ClientID + " " + p.ExpirationPolicy
}

type ProviderPolicy struct {
        rows []ProvidersRow
}

func (p *ProviderPolicy) AddRow(row ProvidersRow) {
        p.rows = append(p.rows, row)
}

func (p *ProviderPolicy) CreateVerifier() (*verifier.Verifier, error) {
        pvs := []verifier.ProviderVerifier{}
        var expirationPolicy verifier.ExpirationPolicy
        var err error
        for _, row := range p.rows {
                var provider verifier.ProviderVerifier
                // TODO: We should handle this issuer matching in a more generic way
                // oidc.local and localhost: are a test issuers
                if row.Issuer == "https://accounts.google.com" ||
                        strings.HasPrefix(row.Issuer, "http://oidc.local") ||
                        strings.HasPrefix(row.Issuer, "http://localhost:") {

                        opts := providers.GetDefaultGoogleOpOptions()
                        opts.Issuer = row.Issuer
                        opts.ClientID = row.ClientID
                        provider = providers.NewGoogleOpWithOptions(opts)
                } else if strings.HasPrefix(row.Issuer, "https://login.microsoftonline.com") {
                        opts := providers.GetDefaultAzureOpOptions()
                        opts.Issuer = row.Issuer
                        opts.ClientID = row.ClientID
                        provider = providers.NewAzureOpWithOptions(opts)
                } else if row.Issuer == "https://gitlab.com" {
                        opts := providers.GetDefaultGitlabOpOptions()
                        opts.Issuer = row.Issuer
                        opts.ClientID = row.ClientID
                        provider = providers.NewGitlabOpWithOptions(opts)
                } else {
                        opts := providers.GetDefaultGoogleOpOptions()
                        opts.Issuer = row.Issuer
                        opts.ClientID = row.ClientID
                        provider = providers.NewGoogleOpWithOptions(opts)
                }

                expirationPolicy, err = row.GetExpirationPolicy()
                if err != nil {
                        return nil, err
                }
                pv := verifier.ProviderVerifierExpires{
                        ProviderVerifier: provider,
                        Expiration:       expirationPolicy,
                }
                pvs = append(pvs, pv)
        }

        if len(pvs) == 0 {
                return nil, fmt.Errorf("no providers configured")
        }
        pktVerifier, err := verifier.NewFromMany(
                pvs,
                verifier.WithExpirationPolicy(expirationPolicy),
        )
        if err != nil {
                return nil, err
        }
        return pktVerifier, nil
}

func (p ProviderPolicy) ToString() string {
        var sb strings.Builder
        for _, row := range p.rows {
                sb.WriteString(row.ToString() + "\n")
        }
        return sb.String()
}

type ProvidersFileLoader struct {
        files.FileLoader
        Path string
}

func NewProviderFileLoader() *ProvidersFileLoader {
        return &ProvidersFileLoader{
                FileLoader: files.FileLoader{
                        Fs:           afero.NewOsFs(),
                        RequiredPerm: files.ModeSystemPerms,
                },
        }
}

func (o *ProvidersFileLoader) LoadProviderPolicy(path string) (*ProviderPolicy, error) {
        content, err := o.FileLoader.LoadFileAtPath(path)
        if err != nil {
                return nil, err
        }
        policy := o.FromTable(content, path)
        return policy, nil
}

// FromTable decodes whitespace delimited input into policy.Policy
func (o ProvidersFileLoader) ToTable(opPolicies ProviderPolicy) files.Table {
        table := files.Table{}
        for _, opPolicy := range opPolicies.rows {
                table.AddRow(opPolicy.Issuer, opPolicy.ClientID, opPolicy.ExpirationPolicy)
        }
        return table
}

// FromTable decodes whitespace delimited input into policy.Policy
// Path is passed only for logging purposes
func (o *ProvidersFileLoader) FromTable(input []byte, path string) *ProviderPolicy {
        table := files.NewTable(input)
        policy := &ProviderPolicy{
                rows: []ProvidersRow{},
        }
        for i, row := range table.GetRows() {
                // Error should not break everyone's ability to login, skip those rows
                if len(row) != 3 {
                        configProblem := files.ConfigProblem{
                                Filepath:            path,
                                OffendingLine:       strings.Join(row, " "),
                                OffendingLineNumber: i,
                                ErrorMessage:        fmt.Sprintf("wrong number of arguments (expected=3, got=%d)", len(row)),
                                Source:              "providers policy file",
                        }
                        files.ConfigProblems().RecordProblem(configProblem)
                        continue
                }
                policyRow := ProvidersRow{
                        Issuer:           row[0],
                        ClientID:         row[1],
                        ExpirationPolicy: row[2], //TODO: Validate this so that we can determine the line number that has the error
                }
                policy.AddRow(policyRow)
        }
        return policy
}

// 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 sshcert

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

        "github.com/lestrrat-go/jwx/v2/jwk"
        "github.com/openpubkey/openpubkey/pktoken"
        "github.com/openpubkey/openpubkey/verifier"
        "golang.org/x/crypto/ssh"
)

type SshCertSmuggler struct {
        SshCert *ssh.Certificate
}

func New(pkt *pktoken.PKToken, principals []string) (*SshCertSmuggler, error) {

        // TODO: assumes email exists in ID Token,
        // this will break for OPs like Azure that do not have email as a claim
        var claims struct {
                Email string `json:"email"`
        }
        if err := json.Unmarshal(pkt.Payload, &claims); err != nil {
                return nil, err
        }

        pubkeySsh, err := sshPubkeyFromPKT(pkt)
        if err != nil {
                return nil, err
        }
        pktCom, err := pkt.Compact()
        if err != nil {
                return nil, err
        }
        sshSmuggler := SshCertSmuggler{
                SshCert: &ssh.Certificate{
                        Key:             pubkeySsh,
                        CertType:        ssh.UserCert,
                        KeyId:           claims.Email,
                        ValidPrincipals: principals,
                        ValidBefore:     ssh.CertTimeInfinity,
                        Permissions: ssh.Permissions{
                                Extensions: map[string]string{
                                        "permit-X11-forwarding":   "",
                                        "permit-agent-forwarding": "",
                                        "permit-port-forwarding":  "",
                                        "permit-pty":              "",
                                        "permit-user-rc":          "",
                                        "openpubkey-pkt":          string(pktCom),
                                },
                        },
                },
        }
        return &sshSmuggler, nil
}

func NewFromAuthorizedKey(certType string, certB64 string) (*SshCertSmuggler, error) {
        if certPubkey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(certType + " " + certB64)); err != nil {
                return nil, err
        } else {
                sshCert, ok := certPubkey.(*ssh.Certificate)
                if !ok {
                        return nil, fmt.Errorf("parsed SSH authorized_key is not an SSH certificate")
                }
                opkcert := &SshCertSmuggler{
                        SshCert: sshCert,
                }
                return opkcert, nil
        }
}

func (s *SshCertSmuggler) SignCert(signerMas ssh.MultiAlgorithmSigner) (*ssh.Certificate, error) {
        if err := s.SshCert.SignCert(rand.Reader, signerMas); err != nil {
                return nil, err
        }
        return s.SshCert, nil
}

func (s *SshCertSmuggler) VerifyCaSig(caPubkey ssh.PublicKey) error {
        certCopy := *(s.SshCert)
        certCopy.Signature = nil
        certBytes := certCopy.Marshal()
        certBytes = certBytes[:len(certBytes)-4] // Drops signature length bytes (see crypto.ssh.certs.go)
        return caPubkey.Verify(certBytes, s.SshCert.Signature)
}

func (s *SshCertSmuggler) GetPKToken() (*pktoken.PKToken, error) {
        pktCom, ok := s.SshCert.Extensions["openpubkey-pkt"]
        if !ok {
                return nil, fmt.Errorf("cert is missing required openpubkey-pkt extension")
        }
        pkt, err := pktoken.NewFromCompact([]byte(pktCom))
        if err != nil {
                return nil, fmt.Errorf("openpubkey-pkt extension in cert failed deserialization: %w", err)
        }
        return pkt, nil
}

func (s *SshCertSmuggler) VerifySshPktCert(ctx context.Context, pktVerifier verifier.Verifier) (*pktoken.PKToken, error) {
        pkt, err := s.GetPKToken()
        if err != nil {
                return nil, fmt.Errorf("openpubkey-pkt extension in cert failed deserialization: %w", err)
        }

        ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
        defer cancel()
        err = pktVerifier.VerifyPKToken(ctxWithTimeout, pkt)
        if err != nil {
                return nil, err
        }

        cic, err := pkt.GetCicValues()
        if err != nil {
                return nil, err
        }
        upk := cic.PublicKey()

        cryptoCertKey := (s.SshCert.Key.(ssh.CryptoPublicKey)).CryptoPublicKey()
        jwkCertKey, err := jwk.FromRaw(cryptoCertKey)
        if err != nil {
                return nil, err
        }

        if jwk.Equal(jwkCertKey, upk) {
                return pkt, nil
        } else {
                return nil, fmt.Errorf("public key 'upk' in PK Token does not match public key in certificate")
        }
}

func sshPubkeyFromPKT(pkt *pktoken.PKToken) (ssh.PublicKey, error) {
        cic, err := pkt.GetCicValues()
        if err != nil {
                return nil, err
        }
        upk := cic.PublicKey()

        var rawkey any
        if err := upk.Raw(&rawkey); err != nil {
                return nil, err
        }
        return ssh.NewPublicKey(rawkey)
}