//  Copyright 2024 Google LLC
//
//  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.

//go:build linux

package oslogin

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"github.com/GoogleCloudPlatform/google-guest-agent/internal/cfg"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/daemon"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/events"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/metadata"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/osinfo"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/pipewatcher"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/run"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/textconfig"
	"github.com/GoogleCloudPlatform/google-guest-agent/internal/utils/file"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

func TestNewModule(t *testing.T) {
	mod := NewModule(context.Background())
	if mod == nil {
		t.Fatalf("NewModule() = nil, want non-nil")
	}

	if mod.ID != osloginModuleID {
		t.Errorf("NewModule().ID = %q, want %q", mod.ID, osloginModuleID)
	}

	if mod.Setup == nil {
		t.Errorf("NewModule().Setup = nil, want non-nil")
	}
}

func TestModuleSetupInputValidity(t *testing.T) {
	mdsJSON := `
	{
		"instance":  {
		}
	}`

	desc, err := metadata.UnmarshalDescriptor(mdsJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", mdsJSON, err)
	}

	tests := []struct {
		name       string
		desc       any
		shouldFail bool
	}{
		{
			name:       "wrong-data",
			desc:       "wrong data",
			shouldFail: true,
		},
		{
			name:       "empty-mds",
			desc:       desc,
			shouldFail: false,
		},
	}

	ctx := context.Background()

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			mod := &osloginModule{}

			err := mod.moduleSetup(ctx, tc.desc)
			if tc.shouldFail != (err != nil) {
				t.Errorf("moduleSetup() = %v, want %v", err, tc.shouldFail)
			}

			// Double call should not fail.
			err = mod.moduleSetup(ctx, tc.desc)
			if tc.shouldFail != (err != nil) {
				t.Errorf("moduleSetup() = %v, want %v", err, tc.shouldFail)
			}
		})
	}
}

func TestMetadataSubscriberInputValidity(t *testing.T) {
	mdsJSON := `
	{
		"instance":  {
		}
	}`

	desc, err := metadata.UnmarshalDescriptor(mdsJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", mdsJSON, err)
	}

	tests := []struct {
		name    string
		desc    any
		err     error
		want    bool
		wantErr bool
	}{
		{
			name:    "wrong-data",
			desc:    "wrong data",
			want:    false,
			wantErr: true,
		},
		{
			name:    "empty-mds",
			desc:    desc,
			want:    true,
			wantErr: false,
		},
		{
			name:    "error-evdata",
			desc:    desc,
			want:    true,
			err:     errors.New("error"),
			wantErr: true,
		},
	}

	ctx := context.Background()

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			mod := &osloginModule{}
			evData := &events.EventData{Data: tc.desc, Error: tc.err}
			got, noop, err := mod.metadataSubscriber(ctx, "evType", nil, evData)
			if (err != nil) != tc.wantErr {
				t.Fatalf("metadataSubscriber() = %v, want error: %t", err, tc.wantErr)
			}
			if !noop {
				t.Errorf("metadataSubscriber() = %v, want: %t", noop, true)
			}
			if got != tc.want {
				t.Errorf("metadataSubscriber() = %v, want %v", got, tc.want)
			}
		})
	}
}

func TestMetadataChanged(t *testing.T) {
	tests := []struct {
		name        string
		prevMDSJSON string
		newMDSJSON  string
		want        bool
	}{
		{
			name: "no-change-basic-mds",
			prevMDSJSON: `
			{
				"instance":  {
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
				}
			}`,
			want: false,
		},
		{
			name: "same-instance-enabled",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-instance-enabled",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-project-enabled",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-project-enabled",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-instance-2fa",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-instance-2fa",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-2fa": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-project-2fa",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-project-2fa",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-2fa": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-2fa": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-instance-sk",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-instance-sk",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-sk": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-project-sk",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-project-sk",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-sk": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-sk": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-instance-certificates",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-instance-certificates",
			prevMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-certificates": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"instance":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			want: true,
		},
		{
			name: "same-project-certificates",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			want: false,
		},
		{
			name: "transition-project-certificates",
			prevMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-certificates": "false"
					}
				}
			}`,
			newMDSJSON: `
			{
				"project":  {
					"attributes": {
						"enable-oslogin-certificates": "true"
					}
				}
			}`,
			want: true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			prevDesc, err := metadata.UnmarshalDescriptor(tc.prevMDSJSON)
			if err != nil {
				t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", tc.prevMDSJSON, err)
			}
			newDesc, err := metadata.UnmarshalDescriptor(tc.newMDSJSON)
			if err != nil {
				t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", tc.newMDSJSON, err)
			}

			mod := &osloginModule{prevMetadata: prevDesc}
			got := mod.metadataChanged(newDesc)
			if got != tc.want {
				t.Errorf("metadataChanged(%v) = %t, want %t", newDesc, got, tc.want)
			}

		})
	}
}

func TestOSLoginSetupError(t *testing.T) {
	mod := &osloginModule{}
	mod.enabled.Store(true)
	mod.pipeEventHandler = &PipeEventHandler{}
	mod.pipeEventWatcher = &pipewatcher.Handle{}
	disableMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "false"
			}
		}
	}`
	disableDesc, err := metadata.UnmarshalDescriptor(disableMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", disableMDSJSON, err)
	}

	got, noop, err := mod.osloginSetup(context.Background(), disableDesc)
	if err == nil {
		t.Errorf("osloginSetup(ctx, %v) = %v, want error", disableDesc, got)
	}
	if noop {
		t.Errorf("osloginSetup(ctx, %v) = %t, want noop to be false", disableDesc, noop)
	}
	if !got {
		t.Errorf("osloginSetup(ctx, %v) = %t, want %t", disableDesc, got, true)
	}
	if !mod.failedConfiguration.Load() {
		t.Errorf("mod.failedConfiguration.Load() = %t, want %t", mod.failedConfiguration.Load(), true)
	}
}

func TestEnableDisable(t *testing.T) {
	// Initialize cfg.
	if err := cfg.Load(nil); err != nil {
		t.Fatalf("failed to load config: %v", err)
	}

	// Create test files and setup test runner.
	module := createTestModule(t)
	createTestFiles(t, module, osloginTestFileOpts{
		testSSHD:     true,
		testNSSwitch: true,
		testPAM:      true,
		testGroup:    true,
	})
	_ = setupTestRunner(t, false)

	enabledMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "true"
			}
		}
	}`
	enabledDesc, err := metadata.UnmarshalDescriptor(enabledMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", enabledMDSJSON, err)
	}

	ctx := context.Background()

	got, noop, err := module.osloginSetup(ctx, enabledDesc)
	if err != nil {
		t.Fatalf("osloginSetup(ctx, %v) = %v, want nil", enabledDesc, err)
	}
	if noop {
		t.Errorf("osloginSetup(ctx, %v) = %t, want noop to be false", enabledDesc, noop)
	}
	if got != true {
		t.Errorf("osloginSetup(ctx, %v) = %t, want %t", enabledDesc, got, true)
	}

	evManager := events.FetchManager()
	if !evManager.IsSubscribed(sshcaPipeWatcherOpts.ReadEventID, pipeWatcherSubscriberID) {
		t.Errorf("pipewatcher.ReadEvent is not subscribed to pipeWatcherSubscriberID, it should be")
	}

	if module.enabled.Load() != true {
		t.Errorf("mod.enabled.Load() = %v, want %v", module.enabled.Load(), true)
	}

	disabledMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "false"
			}
		}
	}`

	disabledDesc, err := metadata.UnmarshalDescriptor(disabledMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", disabledMDSJSON, err)
	}

	got, noop, err = module.osloginSetup(ctx, disabledDesc)
	if err != nil {
		t.Fatalf("osloginSetup(ctx, %v) = %v, want nil", disabledDesc, err)
	}
	if noop {
		t.Errorf("osloginSetup(ctx, %v) = %t, want noop to be false", disabledDesc, noop)
	}
	if !got {
		t.Errorf("osloginSetup(ctx, %v) = %t, want %t", disabledDesc, got, true)
	}

	if evManager.IsSubscribed(sshcaPipeWatcherOpts.ReadEventID, pipeWatcherSubscriberID) {
		t.Errorf("pipewatcher.ReadEvent is still subscribed to pipeWatcherSubscriberID, it should not be")
	}

	if module.enabled.Load() {
		t.Errorf("mod.enabled.Load() = %v, want %v", module.enabled.Load(), false)
	}
}

func TestDisableOSLoginErrors(t *testing.T) {
	if err := cfg.Load(nil); err != nil {
		t.Fatalf("Failed to load config: %v", err)
	}

	tests := []struct {
		name      string
		fileOpts  osloginTestFileOpts
		runnerErr bool
		wantErr   bool
	}{
		{
			name:    "fail-sshd-cleanup",
			wantErr: true,
		},
		// Non existent nss file should not cause an error.
		{
			name:     "fail-nss-cleanup",
			fileOpts: osloginTestFileOpts{testSSHD: true},
			wantErr:  true,
		},
		// Non existent pam file should not cause an error.
		{
			name:     "fail-pam-cleanup",
			fileOpts: osloginTestFileOpts{testSSHD: true, testNSSwitch: true},
		},
		// Non existent group file should not cause an error.
		{
			name:     "fail-group-cleanup",
			fileOpts: osloginTestFileOpts{testSSHD: true, testNSSwitch: true, testPAM: true},
		},
		{
			name:      "fail-restart-services",
			fileOpts:  osloginTestFileOpts{testSSHD: true, testNSSwitch: true, testPAM: true, testGroup: true},
			runnerErr: true,
			wantErr:   true,
		},
	}

	// Mock oslogin enabled metadata.
	enabledMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "true"
			}
		}
	}`
	enabledDesc, err := metadata.UnmarshalDescriptor(enabledMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", enabledMDSJSON, err)
	}

	// Run tests.
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			module := createTestModule(t)
			module.services = map[daemon.RestartMethod][]serviceRestartConfig{
				daemon.TryRestart: {
					{
						services: []string{"service1"},
					},
				},
			}
			createTestFiles(t, module, test.fileOpts)
			_ = setupTestRunner(t, test.runnerErr)

			module.osloginSetup(context.Background(), enabledDesc)

			// Now test for errors.
			if err := module.disableOSLogin(context.Background(), events.FetchManager()); test.wantErr != (err != nil) {
				t.Errorf("disableOSLogin(ctx, evManager) = %v, want err: %t", err, test.wantErr)
			}
		})
	}
}

func TestSetupOpenSSH(t *testing.T) {
	if err := cfg.Load(nil); err != nil {
		t.Fatalf("Failed to load config: %v", err)
	}
	module := createTestModule(t)
	createTestFiles(t, module, osloginTestFileOpts{testSSHD: true})

	tests := []struct {
		name            string
		desc            string
		certAuthEnabled bool
		expectedLines   []string
	}{
		{
			name: "no_cert_no_2fa",
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}
			`,
			certAuthEnabled: true,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "cert",
			certAuthEnabled: true,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-certificates": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "2fa",
			certAuthEnabled: true,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-2fa": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"AuthenticationMethods publickey,keyboard-interactive",
				"ChallengeResponseAuthentication yes",
				"Include /var/google-users.d/*",
				"Match User sa_*",
				"AuthenticationMethods publickey",
			},
		},
		{
			name:            "cert_and_2fa",
			certAuthEnabled: true,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-2fa": "true",
						"enable-oslogin-certificates": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				"AuthenticationMethods publickey,keyboard-interactive",
				"ChallengeResponseAuthentication yes",
				"Include /var/google-users.d/*",
				"Match User sa_*",
				"AuthenticationMethods publickey",
			},
		},
		{
			name:            "sk",
			certAuthEnabled: true,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-sk": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandSKPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "cert_disabled",
			certAuthEnabled: false,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "cert_metadata_disabled_cert_auth_enabled",
			certAuthEnabled: true,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-certificates": "false"
					}
				}
			}
			`,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "cert_metadata_enabled_cert_auth_disabled",
			certAuthEnabled: false,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-certificates": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
		{
			name:            "2fa_cert_disabled",
			certAuthEnabled: false,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-2fa": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"AuthenticationMethods publickey,keyboard-interactive",
				"ChallengeResponseAuthentication yes",
				"Include /var/google-users.d/*",
				"Match User sa_*",
				"AuthenticationMethods publickey",
			},
		},
		{
			name:            "sk_cert_disabled",
			certAuthEnabled: false,
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-sk": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandSKPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			desc, err := metadata.UnmarshalDescriptor(test.desc)
			if err != nil {
				t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", test.desc, err)
			}

			if err := cfg.Load(nil); err != nil {
				t.Fatalf("Failed to load config: %v", err)
			}
			cfg.Retrieve().OSLogin.CertAuthentication = test.certAuthEnabled

			err = module.setupOpenSSH(desc)
			if err != nil {
				t.Fatalf("setupOpenSSH(desc) = %v, want nil", err)
			}

			// Check that the file has expected contents.
			checkTestFile(t, module.sshdConfigPath, test.expectedLines)
		})
	}
}

func TestSetupNSSwitch(t *testing.T) {
	module := createTestModule(t)
	createTestFiles(t, module, osloginTestFileOpts{testNSSwitch: true})

	if err := module.setupNSSwitch(false); err != nil {
		t.Fatalf("setupNSSwitch(false) = %v, want nil", err)
	}

	checkTestFile(t, module.nsswitchConfigPath, []string{
		"passwd: files cache_oslogin oslogin",
		"group: files cache_oslogin oslogin",
	})
}

func TestSetupPAM(t *testing.T) {
	tests := []struct {
		name          string
		desc          string
		expectedLines []string
	}{
		{
			name: "default",
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"auth [default=ignore] pam_group.so",
				"session [success=ok default=ignore] pam_mkhomedir.so",
			},
		},
		{
			name: "two-factor",
			desc: `
			{
				"instance": {
					"attributes": {
						"enable-oslogin": "true",
						"enable-oslogin-2fa": "true"
					}
				}
			}
			`,
			expectedLines: []string{
				"auth [success=done perm_denied=die default=ignore] pam_oslogin_login.so",
				"auth [default=ignore] pam_group.so",
				"session [success=ok default=ignore] pam_mkhomedir.so",
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			module := createTestModule(t)
			desc, err := metadata.UnmarshalDescriptor(test.desc)
			if err != nil {
				t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", test.desc, err)
			}
			createTestFiles(t, module, osloginTestFileOpts{testPAM: true})

			if err := module.setupPAM(desc); err != nil {
				t.Fatalf("setupPAM() = %v, want nil", err)
			}

			checkTestFile(t, module.pamConfigPath, test.expectedLines)
		})
	}
}

func TestSetupGroup(t *testing.T) {
	module := createTestModule(t)
	createTestFiles(t, module, osloginTestFileOpts{testGroup: true})

	if err := module.setupGroup(); err != nil {
		t.Fatalf("setupGroup() = %v, want nil", err)
	}

	checkTestFile(t, module.groupConfigPath, []string{
		"sshd;*;*;Al0000-2400;video",
	})
}

func TestRestartServices(t *testing.T) {
	unknownMethod := daemon.RestartMethod(50)

	tests := []struct {
		name            string
		returnErr       bool
		services        map[daemon.RestartMethod][]serviceRestartConfig
		expectedCommand string
		expectedArgs    []string
	}{
		{
			name: "restart",
			services: map[daemon.RestartMethod][]serviceRestartConfig{
				daemon.TryRestart: {
					{
						services: []string{"service1"},
					},
				},
			},
			expectedCommand: "systemctl",
			expectedArgs:    []string{"try-restart", "service1"},
		},
		{
			name: "reload_restart",
			services: map[daemon.RestartMethod][]serviceRestartConfig{
				daemon.ReloadOrRestart: {
					{
						services: []string{"service1"},
					},
				},
			},
			expectedCommand: "systemctl",
			expectedArgs:    []string{"reload-or-restart", "service1"},
		},
		{
			name:      "error",
			returnErr: true,
			services: map[daemon.RestartMethod][]serviceRestartConfig{
				daemon.TryRestart: {
					{
						services: []string{"service1"},
					},
				},
			},
		},
		{
			name:      "unknown_method",
			returnErr: true,
			services: map[daemon.RestartMethod][]serviceRestartConfig{
				unknownMethod: {
					{
						services: []string{"service1"},
					},
				},
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			runner := setupTestRunner(t, test.returnErr)
			module := &osloginModule{
				services: test.services,
			}

			if err := module.restartServices(context.Background()); (err == nil) == test.returnErr {
				t.Fatalf("restartServices(ctx) = %v, want %v", err, test.returnErr)
			}
			if test.returnErr {
				return
			}

			args, found := runner.seenCommand[test.expectedCommand]
			if !found {
				t.Fatalf("restartServices(ctx) did not call %q", test.expectedCommand)
			}
			if len(args) != 2 {
				t.Fatalf("restartServices(ctx) called %q with %d args, want 2", test.expectedCommand, len(args))
			}
			if diff := cmp.Diff(test.expectedArgs, args[1]); diff != "" {
				t.Fatalf("restartServices(ctx) called %q with args diff (-want +got): %v", test.expectedCommand, diff)
			}
		})
	}
}

func TestSetupOSLoginDirs(t *testing.T) {
	temp := t.TempDir()
	module := &osloginModule{
		osloginDirs: []string{
			filepath.Join(temp, "google-users.d"),
			filepath.Join(temp, "/tmp/google-sudoers.d"),
		},
	}
	execLookPath = func(string) (string, error) { return "restoreconn", nil }
	t.Cleanup(func() { execLookPath = exec.LookPath })
	testRunner := setupTestRunner(t, false)

	if err := module.setupOSLoginDirs(context.Background()); err != nil {
		t.Fatalf("setupOSLoginDirs(ctx) = %v, want nil", err)
	}

	for _, dir := range module.osloginDirs {
		if !file.Exists(dir, file.TypeDir) {
			t.Fatalf("setupOSLoginDirs(ctx) did not create dir %q", dir)
		}
	}

	expectedCommandArgs := [][]string{
		{module.osloginDirs[0]},
		{module.osloginDirs[1]},
	}

	seenCommandArgs, found := testRunner.seenCommand["restoreconn"]
	if !found {
		t.Fatalf("setupOSLoginDirs(ctx) did not call restoreconn")
	}
	if diff := cmp.Diff(expectedCommandArgs, seenCommandArgs); diff != "" {
		t.Fatalf("setupOSLoginDirs(ctx) called restoreconn with diff (-want +got): %v", diff)
	}
}

func TestSetupOSLoginSudoers(t *testing.T) {
	temp := t.TempDir()
	module := &osloginModule{
		sudoers: filepath.Join(temp, "google-oslogin-sudoers"),
	}

	if err := module.setupOSLoginSudoers(); err != nil {
		t.Fatalf("createOSLoginSudoers(ctx) = %v, want nil", err)
	}

	checkTestFile(t, module.sudoers, []string{
		"#include /var/google-sudoers.d",
	})
}

func TestRemoveDeprecatedEntries(t *testing.T) {
	// Create a test file with a deprecated entry.
	temp := t.TempDir()
	testFile := filepath.Join(temp, "test_file")
	testContents := "depKey depEntry\nnotDepKey notDepEntry\n"
	if err := os.WriteFile(testFile, []byte(testContents), osloginConfigMode); err != nil {
		t.Fatalf("failed to setup test file: %v", err)
	}

	module := &osloginModule{
		deprecatedEntries: map[string][]*textconfig.Entry{
			testFile: []*textconfig.Entry{
				textconfig.NewEntry("depKey", "depEntry"),
			},
		},
	}

	if err := module.removeDeprecatedEntries(); err != nil {
		t.Fatalf("removeDeprecatedEntries() = %v, want nil", err)
	}

	// Double check that the file no longer has the deprecated entry.
	checkTestFile(t, testFile, []string{"notDepKey notDepEntry"})
}

func TestSetupUsrEtcOSLoginDirs(t *testing.T) {
	oldSles16Map := sles16Map
	oldOsinfoRead := osinfoRead
	t.Cleanup(func() {
		sles16Map = oldSles16Map
		osinfoRead = oldOsinfoRead
	})

	wantMap := map[string]string{
		"/usr/etc/ssh/sshd_config":     "/etc/ssh/sshd_config",
		"/usr/etc/nsswitch.conf":       "/etc/nsswitch.conf",
		"/usr/lib/pam.d/sshd":          "/etc/pam.d/sshd",
		"/usr/etc/security/group.conf": "/etc/security/group.conf",
	}
	if diff := cmp.Diff(wantMap, sles16Map); diff != "" {
		t.Fatalf("sles16Map unexpected diff (-want +got):\n%s", diff)
	}

	ctx := context.Background()

	tests := []struct {
		name           string
		info           osinfo.OSInfo
		createSrc      bool
		createDst      bool
		dstContent     string
		dstShouldExist bool
		prevSetup      bool
	}{
		{
			name:           "debian12-no-copy",
			info:           osinfo.OSInfo{OS: "debian", Version: osinfo.Ver{Major: 12}},
			createSrc:      true,
			createDst:      false,
			dstShouldExist: false,
		},
		{
			name:           "opensuse15-no-copy",
			info:           osinfo.OSInfo{OS: "opensuse-leap", Version: osinfo.Ver{Major: 15}},
			createSrc:      true,
			createDst:      false,
			dstShouldExist: false,
		},
		{
			name:           "sles15-no-copy",
			info:           osinfo.OSInfo{OS: "sles", Version: osinfo.Ver{Major: 15}},
			createSrc:      true,
			createDst:      false,
			dstShouldExist: false,
		},
		{
			name:           "sles16-copy",
			info:           osinfo.OSInfo{OS: "sles", Version: osinfo.Ver{Major: 16}},
			createSrc:      true,
			createDst:      false,
			dstShouldExist: true,
			dstContent:     "test",
		},
		{
			name:           "opensuse16-copy",
			info:           osinfo.OSInfo{OS: "opensuse-leap", Version: osinfo.Ver{Major: 16}},
			createSrc:      true,
			createDst:      false,
			dstShouldExist: true,
			dstContent:     "test",
		},
		{
			name:           "sles16-no-copy-if-exists",
			info:           osinfo.OSInfo{OS: "sles", Version: osinfo.Ver{Major: 16}},
			createSrc:      false,
			createDst:      true,
			dstContent:     "exists",
			dstShouldExist: true,
		},
		{
			name:           "sles16-no-copy-if-already-setup",
			info:           osinfo.OSInfo{OS: "sles", Version: osinfo.Ver{Major: 16}},
			createSrc:      false,
			createDst:      true,
			dstContent:     "exists",
			dstShouldExist: true,
			prevSetup:      true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mod := &osloginModule{}
			mod.sles16Setup.Store(tt.prevSetup)

			usrDir := t.TempDir()
			etcDir := t.TempDir()
			src := filepath.Join(usrDir, "nsswitch.conf")
			dst := filepath.Join(etcDir, "nsswitch.conf")
			sles16Map = map[string]string{
				src: dst,
			}

			if err := os.MkdirAll(filepath.Dir(src), 0755); err != nil {
				t.Fatalf("Failed to create dir for %s: %v", src, err)
			}
			if tt.createSrc {
				if err := os.WriteFile(src, []byte("test"), 0644); err != nil {
					t.Fatalf("Failed to write to %s: %v", src, err)
				}
			}

			if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil {
				t.Fatalf("Failed to create dir for %s: %v", dst, err)
			}
			if tt.createDst {
				if err := os.WriteFile(dst, []byte(tt.dstContent), 0644); err != nil {
					t.Fatalf("Failed to write to %s: %v", dst, err)
				}
			}

			osinfoRead = func() osinfo.OSInfo {
				return tt.info
			}

			if err := mod.setupSles16OSLoginDirs(ctx); err != nil {
				t.Fatalf("setupUsrEtcOSLoginDirs() returned err: %v, want nil", err)
			}

			if got := file.Exists(dst, file.TypeFile); got != tt.dstShouldExist {
				t.Errorf("Destination file %s exists: %t, want: %t", dst, got, tt.dstShouldExist)
			}
			if tt.dstShouldExist {
				got, err := os.ReadFile(dst)
				if err != nil {
					t.Fatalf("Failed to read destination file %s: %v", dst, err)
				}
				if string(got) != tt.dstContent {
					t.Errorf("Destination file %s content changed to %s, want %s", dst, string(got), tt.dstContent)
				}
			}
			if got := mod.sles16Setup.Load(); !got {
				t.Errorf("sles16Setup is %t, want true", got)
			}
		})
	}
}

func TestOSLoginSetupSLES16Failure(t *testing.T) {
	oldSles16Map := sles16Map
	oldOsinfoRead := osinfoRead
	t.Cleanup(func() {
		sles16Map = oldSles16Map
		osinfoRead = oldOsinfoRead
	})

	if err := cfg.Load(nil); err != nil {
		t.Fatalf("Failed to load config: %v", err)
	}

	mod := createTestModule(t)
	mod.enabled.Store(true)
	usrDir := t.TempDir()

	// sles16Map source files are not created, so copy will fail.
	sles16Map = map[string]string{
		filepath.Join(usrDir, "sshd_config"): mod.sshdConfigPath,
	}

	osinfoRead = func() osinfo.OSInfo {
		return osinfo.OSInfo{OS: "sles", Version: osinfo.Ver{Major: 16}}
	}

	enabledMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "true"
			}
		}
	}`
	enabledDesc, err := metadata.UnmarshalDescriptor(enabledMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", enabledMDSJSON, err)
	}

	ctx := context.Background()

	shouldContinue, noop, err := mod.osloginSetup(ctx, enabledDesc)
	if err == nil {
		t.Fatalf("osloginSetup(ctx, %+v) succeeded, want error", enabledDesc)
	}

	wantErr := "failed to setup OSLogin directories"
	if !strings.Contains(err.Error(), wantErr) {
		t.Fatalf("osloginSetup(ctx, %+v) error is %q, want error to contain %q", enabledDesc, err.Error(), wantErr)
	}

	if !mod.failedConfiguration.Load() {
		t.Errorf("osloginSetup(ctx, %+v) failedConfiguration is false, want true", enabledDesc)
	}
	if mod.enabled.Load() {
		t.Errorf("osloginSetup(ctx, %+v) enabled is true, want false", enabledDesc)
	}
	if !shouldContinue {
		t.Errorf("osloginSetup(ctx, %+v) shouldContinue is %t, want true", enabledDesc, shouldContinue)
	}
	if noop {
		t.Errorf("osloginSetup(ctx, %+v) noop is %t, want false", enabledDesc, noop)
	}
	if mod.sles16Setup.Load() {
		t.Errorf("osloginSetup(ctx, %+v) sles16Setup is true, want false", enabledDesc)
	}
}

func TestRetryFailConfiguration(t *testing.T) {
	if err := cfg.Load(nil); err != nil {
		t.Fatalf("failed to load config: %v", err)
	}

	tests := []struct {
		name      string
		services  map[daemon.RestartMethod][]serviceRestartConfig
		fileOpts  osloginTestFileOpts
		runnerErr bool
	}{
		{
			name: "retry-sshd",
		},
		{
			name: "retry-nss",
			fileOpts: osloginTestFileOpts{
				testSSHD: true,
			},
		},
		{
			name: "retry-restart-services",
			services: map[daemon.RestartMethod][]serviceRestartConfig{
				daemon.TryRestart: {
					{
						protocol: serviceRestartOptional,
						services: []string{"service1"},
					},
				},
			},
			fileOpts: osloginTestFileOpts{
				testSSHD:     true,
				testNSSwitch: true,
				testPAM:      true,
				testGroup:    true,
			},
			runnerErr: true,
		},
		{
			name: "retry-nss-cache-fill",
			fileOpts: osloginTestFileOpts{
				testSSHD:     true,
				testNSSwitch: true,
				testPAM:      true,
				testGroup:    true,
			},
			runnerErr: true,
		},
	}

	// Mock oslogin enabled metadata.
	enabledMDSJSON := `
	{
		"instance":  {
			"attributes": {
				"enable-oslogin": "true"
			}
		}
	}`
	enabledDesc, err := metadata.UnmarshalDescriptor(enabledMDSJSON)
	if err != nil {
		t.Fatalf("metadata.UnmarshalDescriptor(%q) = %v, want nil", enabledMDSJSON, err)
	}

	// Run tests.
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			module := createTestModule(t)
			module.services = test.services
			createTestFiles(t, module, test.fileOpts)
			_ = setupTestRunner(t, test.runnerErr)

			ok, noop, err := module.osloginSetup(context.Background(), enabledDesc)
			if err == nil {
				t.Fatalf("osloginSetup(ctx, %v) = nil, want error", enabledDesc)
			}
			if noop {
				t.Fatalf("osloginSetup(ctx, %v) = %t, want false", enabledDesc, noop)
			}
			if !ok {
				t.Fatalf("osloginSetup(ctx, %v) = %t, want %t", enabledDesc, ok, true)
			}

			// Initial setup somehow succeeded.
			if !module.failedConfiguration.Load() {
				t.Fatalf("osloginSetup(ctx, %v) did not set failedConfiguration", enabledDesc)
			}

			// Now make sure things can pass.
			_ = setupTestRunner(t, false)
			createTestFiles(t, module, osloginTestFileOpts{
				testSSHD:     true,
				testNSSwitch: true,
				testPAM:      true,
				testGroup:    true,
			})

			// Retry should succeed.
			ok, noop, err = module.osloginSetup(context.Background(), enabledDesc)
			if err != nil {
				t.Fatalf("osloginSetup(ctx, %v) = %v, want nil", enabledDesc, err)
			}
			if noop {
				t.Fatalf("osloginSetup(ctx, %v) = %t, want false", enabledDesc, noop)
			}
			if !ok {
				t.Fatalf("osloginSetup(ctx, %v) = %t, want %t", enabledDesc, ok, true)
			}

			if module.failedConfiguration.Load() {
				t.Fatalf("osloginSetup(ctx, %v) did not clear failedConfiguration", enabledDesc)
			}

			// Check that the files do not have duplicate blocks.
			checkTestFile(t, module.sshdConfigPath, []string{
				"TrustedUserCAKeys /etc/ssh/oslogin_trustedca.pub",
				"AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k",
				"AuthorizedPrincipalsCommandUser root",
				fmt.Sprintf("AuthorizedKeysCommand %s", module.authorizedKeysCommandPaths[0]),
				"AuthorizedKeysCommandUser root",
				"Include /var/google-users.d/*",
			})
			checkTestFile(t, module.nsswitchConfigPath, []string{
				"passwd: files cache_oslogin oslogin",
				"group: files cache_oslogin oslogin",
			})
			checkTestFile(t, module.pamConfigPath, []string{
				"auth [default=ignore] pam_group.so",
				"session [success=ok default=ignore] pam_mkhomedir.so",
			})
			checkTestFile(t, module.groupConfigPath, []string{
				"sshd;*;*;Al0000-2400;video",
			})
		})
	}
}

type testRunner struct {
	returnErr   bool
	seenCommand map[string][][]string
}

func (t *testRunner) WithContext(ctx context.Context, opts run.Options) (*run.Result, error) {
	t.seenCommand[opts.Name] = append(t.seenCommand[opts.Name], opts.Args)

	if t.returnErr {
		return nil, errors.New("error")
	}
	return &run.Result{Output: "1 loaded units listed."}, nil
}

func setupTestRunner(t *testing.T, returnErr bool) *testRunner {
	testRunner := &testRunner{
		returnErr:   returnErr,
		seenCommand: make(map[string][][]string),
	}

	oldClient := run.Client
	run.Client = testRunner
	t.Cleanup(func() {
		run.Client = oldClient
	})
	return testRunner
}

type osloginTestFileOpts struct {
	testSSHD     bool
	testNSSwitch bool
	testPAM      bool
	testGroup    bool
}

func createTestModule(t *testing.T) *osloginModule {
	temp := t.TempDir()
	return &osloginModule{
		sshdConfigPath:               filepath.Join(temp, "sshd_config"),
		nsswitchConfigPath:           filepath.Join(temp, "nss_switch.conf"),
		pamConfigPath:                filepath.Join(temp, "sshd"),
		groupConfigPath:              filepath.Join(temp, "group.conf"),
		authorizedKeysCommandPaths:   []string{filepath.Join(temp, "google_authorized_keys")},
		authorizedKeysCommandSKPaths: []string{filepath.Join(temp, "google_authorized_keys_sk")},
		osloginDirs:                  []string{filepath.Join(temp, "google-users.d"), filepath.Join(temp, "google-sudoers.d")},
		sudoers:                      filepath.Join(temp, "google-oslogin-sudoers"),
		deprecatedEntries:            make(map[string][]*textconfig.Entry),
	}
}

func createTestFiles(t *testing.T, module *osloginModule, opts osloginTestFileOpts) {
	if opts.testSSHD {
		if err := os.WriteFile(module.sshdConfigPath, nil, osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test sshd_config file: %v", err)
		}

		// Create test files for AuthorizedKeysCommand and AuthorizedKeysCommandSK.
		if err := os.WriteFile(module.authorizedKeysCommandPaths[0], nil, osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test authorized_keys_command file: %v", err)
		}
		if err := os.WriteFile(module.authorizedKeysCommandSKPaths[0], nil, osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test authorized_keys_command_sk file: %v", err)
		}
	}
	if opts.testNSSwitch {
		if err := os.WriteFile(module.nsswitchConfigPath, []byte("passwd: files\ngroup: files"), osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test nsswitch.conf file: %v", err)
		}
	}
	if opts.testPAM {
		if err := os.WriteFile(module.pamConfigPath, nil, osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test pam.conf file: %v", err)
		}
	}
	if opts.testGroup {
		if err := os.WriteFile(module.groupConfigPath, nil, osloginConfigMode); err != nil {
			t.Fatalf("failed to setup test group.conf file: %v", err)
		}
	}
}

// checkTestFile checks that the provided file has the expected lines.
// It ignores any lines that are commented out or empty.
func checkTestFile(t *testing.T, path string, expectedLines []string) {
	t.Helper()

	contents, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("failed to read test file: %v", err)
	}
	lines := strings.Split(string(contents), "\n")

	if diff := cmp.Diff(expectedLines, lines, cmpopts.IgnoreSliceElements(func(e string) bool {
		return e == "" || strings.HasPrefix(e, "#")
	})); diff != "" {
		t.Errorf("test file %q written with diff (-want +got): %v", filepath.Base(path), diff)
	}
}
