package main

import (
	"context"
	"net"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/shellhub-io/shellhub/pkg/api/authorizer"
	"github.com/shellhub-io/shellhub/pkg/api/requests"
	"github.com/shellhub-io/shellhub/pkg/models"
	"github.com/shellhub-io/shellhub/tests/environment"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

const approvalWait = 60 * time.Second

func dialSSH(ctx context.Context, addr, sshid string, signers []ssh.Signer, prompts chan<- string, answers <-chan string) error {
	auth := []ssh.AuthMethod{ssh.PublicKeysCallback(func() ([]ssh.Signer, error) { return signers, nil })}

	if prompts != nil {
		auth = append(auth, ssh.KeyboardInteractive(func(_, instruction string, questions []string, _ []bool) ([]string, error) {
			select {
			case prompts <- instruction:
			default:
			}

			if len(questions) == 0 {
				return nil, nil
			}

			if answers == nil {
				return make([]string, len(questions)), nil
			}

			select {
			case answer := <-answers:
				return []string{answer}, nil
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}))
	}

	return handshake(ctx, addr, sshid, auth)
}

func handshake(ctx context.Context, addr, sshid string, auth []ssh.AuthMethod) error {
	return handshakeWith(ctx, addr, &ssh.ClientConfig{
		User:            sshid,
		Auth:            auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // a test dialing its own throwaway server
	})
}

func handshakeWith(ctx context.Context, addr string, config *ssh.ClientConfig) error {
	dialer := new(net.Dialer)

	conn, err := dialer.DialContext(ctx, "tcp", addr)
	if err != nil {
		return err
	}

	stop := context.AfterFunc(ctx, func() { _ = conn.Close() })
	defer stop()

	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		_ = conn.Close()

		return err
	}

	_ = ssh.NewClient(c, chans, reqs).Close()

	return nil
}

func enrollIdentity(ctx context.Context, t *testing.T, compose *environment.DockerCompose, publicKey string) {
	t.Helper()

	resp, err := compose.R(ctx).
		SetBody(map[string]string{"name": "enrolled", "data": publicKey, "source": "manual"}).
		Post("/api/ssh-identities")
	require.NoError(t, err)
	require.Contains(t, []int{200, 201}, resp.StatusCode(), "failed to enroll the identity: %s", resp.String())
}

func requireLoginRefused(t *testing.T, compose *environment.DockerCompose, device *models.Device, signer ssh.Signer, wantServerLogs ...string) {
	t.Helper()

	err := dialSSH(t.Context(), compose.SSHAddress(), deviceSSHID(device), []ssh.Signer{signer}, nil, nil)
	require.Error(t, err)

	for _, substr := range wantServerLogs {
		compose.AwaitServerLog(t, substr)
	}
}

var approvalCodePattern = regexp.MustCompile(`/ssh-identities/new/([2-9A-Z]{8})`)

// TestIdentityAccessPolicy covers what decides a login in a new namespace:
// an unknown key is held for approval and let in once confirmed, default-deny leaves a member
// without a grant outside, a deny beats the allow that would otherwise let the owner in, and a
// namespace stripped of every policy admits nobody. Password refusal and a pre-enrolled identity
// connecting are covered by [TestSSHNewNamespace].
func TestIdentityAccessPolicy(t *testing.T) {
	t.Parallel()

	t.Run("an unknown key is held for approval and let in once confirmed", func(t *testing.T) {
		ctx := context.Background()
		compose := newSSHEnvironment(t, ctx, newNamespace)
		_, device := startAcceptedAgent(t, ctx, compose)

		signer, _ := newSigner(t)

		prompts := make(chan string, 8)
		answers := make(chan string, 1)
		dialed := make(chan error, 1)

		go func() {
			dialed <- dialSSH(ctx, compose.SSHAddress(), deviceSSHID(device), []ssh.Signer{signer}, prompts, answers)
		}()

		var code string

		deadline := time.After(approvalWait)

		for code == "" {
			select {
			case instruction := <-prompts:
				if match := approvalCodePattern.FindStringSubmatch(instruction); match != nil {
					code = match[1]
				}
			case err := <-dialed:
				require.NoError(t, err, "the login failed before an approval code was offered")
				require.Fail(t, "the login completed without an approval code being offered")
			case <-deadline:
				require.Fail(t, "the gateway never offered an approval code")
			}
		}

		confirmation := new(models.SSHApprovalConfirmation)
		resp, err := compose.R(ctx).SetResult(confirmation).Post("/api/ssh-approvals/" + code + "/confirm")
		require.NoError(t, err)
		require.Equal(t, 200, resp.StatusCode())
		require.NotEmpty(t, confirmation.ConfirmationCode,
			"confirming has to hand back a code, since that is what the terminal asks for")

		answers <- confirmation.ConfirmationCode

		select {
		case err := <-dialed:
			require.NoError(t, err, "the login should resume once the confirmation code is typed")
		case <-time.After(approvalWait):
			require.Fail(t, "the login never resumed after the approval was confirmed")
		}

		identities := []models.SSHIdentity{}

		resp, err = compose.R(ctx).SetResult(&identities).Get("/api/ssh-identities")
		require.NoError(t, err)
		require.Equal(t, 200, resp.StatusCode())
		require.Len(t, identities, 1)
		assert.Equal(t, models.SSHIdentitySourceApproval, identities[0].Source)
		assert.Equal(t, ssh.FingerprintSHA256(signer.PublicKey()), identities[0].Fingerprint,
			"the identity should hold the key that dialled")
	})

	t.Run("a client that cannot answer a challenge is refused at once", func(t *testing.T) {
		ctx := context.Background()
		compose := newSSHEnvironment(t, ctx, newNamespace)
		_, device := startAcceptedAgent(t, ctx, compose)

		signer, _ := newSigner(t)

		started := time.Now() //nolint:forbidigo // elapsed time of the login, which is what this test asserts on
		err := dialSSH(ctx, compose.SSHAddress(), deviceSSHID(device), []ssh.Signer{signer}, nil, nil)
		elapsed := time.Since(started)

		require.Error(t, err, "an unenrolled key must not get in without an approval")

		assert.Less(t, elapsed, 5*time.Second,
			"a client with no way to answer must be refused, not held while someone is asked")
	})

	t.Run("an agent full of unenrolled keys is asked once and still gets in", func(t *testing.T) {
		ctx := context.Background()
		compose := newSSHEnvironment(t, ctx, newNamespace)
		_, device := startAcceptedAgent(t, ctx, compose)

		enrolled, enrolledKey := newSigner(t)
		enrollIdentity(ctx, t, compose, enrolledKey)

		first, _ := newSigner(t)
		second, _ := newSigner(t)

		prompts := make(chan string, 8)
		dialed := make(chan error, 1)

		go func() {
			dialed <- dialSSH(ctx, compose.SSHAddress(), deviceSSHID(device),
				[]ssh.Signer{first, second, enrolled}, prompts, nil)
		}()

		select {
		case err := <-dialed:
			require.NoError(t, err, "the enrolled key behind the unenrolled ones should still get in")
		case <-time.After(approvalWait):
			require.Fail(t, "the login never finished")
		}

		close(prompts)

		asked := map[string]bool{}
		for instruction := range prompts {
			if strings.Contains(instruction, "doesn't know this SSH key") {
				asked[instruction] = true
			}
		}

		assert.Len(t, asked, 1,
			"one connection is asked about one key, however many unenrolled ones it holds")
	})

	tests := []struct {
		name           string
		refuse         func(t *testing.T, compose *environment.DockerCompose) ssh.Signer
		wantServerLogs []string
	}{
		{
			name: "a member no policy grants is refused",
			refuse: func(t *testing.T, compose *environment.DockerCompose) ssh.Signer {
				t.Helper()

				compose.NewUser(t, "member", "member@ossystems.com.br", ShellHubPassword)
				compose.NewMember(t, "member", ShellHubNamespaceName, string(authorizer.RoleOperator))

				auth := compose.AuthUser(t, "member", ShellHubPassword)
				require.NotNil(t, auth.Tenant)
				require.Equal(t, ShellHubNamespace, *auth.Tenant)

				signer, data := newSigner(t)
				compose.EnrollIdentityAs(t, auth.Token, "member", data)

				return signer
			},
			wantServerLogs: []string{"reason=" + string(models.ReasonNoGrant)},
		},
		{
			name: "a deny policy beats the allow that would grant the owner",
			refuse: func(t *testing.T, compose *environment.DockerCompose) ssh.Signer {
				t.Helper()

				signer, data := newSigner(t)
				compose.EnrollIdentity(t, "owner", data)

				compose.CreateAccessPolicy(t, &requests.AccessPolicyCreate{
					Name:    "no root",
					Subject: requests.AccessPolicySubject{Type: string(models.PolicySubjectAllMembers)},
					Filter:  requests.AccessPolicyFilter{All: true},
					Logins:  []string{ShellHubAgentUsername},
					Action:  string(models.PolicyActionDeny),
				})

				return signer
			},
			wantServerLogs: []string{"reason=" + string(models.ReasonDeniedByPolicy)},
		},
		{
			name: "a namespace with no policy at all refuses everyone",
			refuse: func(t *testing.T, compose *environment.DockerCompose) ssh.Signer {
				t.Helper()

				signer, data := newSigner(t)
				compose.EnrollIdentity(t, "owner", data)

				policies := []models.AccessPolicy{}

				resp, err := compose.R(t.Context()).SetResult(&policies).Get("/api/access-policies")
				require.NoError(t, err)
				require.Equal(t, 200, resp.StatusCode())
				require.NotEmpty(t, policies, "a new namespace is seeded with the owner policy")

				for _, policy := range policies {
					resp, err := compose.R(t.Context()).Delete("/api/access-policies/" + policy.ID)
					require.NoError(t, err)
					require.Equal(t, 200, resp.StatusCode())
				}

				return signer
			},
			wantServerLogs: []string{
				"destination device did not pass the connection evaluation",
				`error="ssh access denied by policy"`,
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			ctx := context.Background()
			compose := newSSHEnvironment(t, ctx, newNamespace)
			_, device := startAcceptedAgent(t, ctx, compose)

			signer := tc.refuse(t, compose)

			requireLoginRefused(t, compose, device, signer, tc.wantServerLogs...)
		})
	}
}