package validator

import (
	"context"
	"fmt"
	"testing"

	"github.com/pkg/errors"

	"code.justin.tv/eventbus/controlplane/internal/policy"
	"code.justin.tv/eventbus/controlplane/internal/validator/mocks"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type ctxid string

func testCtx(ctx context.Context, id string) context.Context {
	return context.WithValue(ctx, ctxid(id), ctxid(id))
}

func TestPublishPermissions(t *testing.T) {
	awsAcctID := "123456789012"
	queueARN := "arn:aws:sqs:us-east-2:444455556666:eventbus-test"
	sourceARNValid := fmt.Sprintf("arn:aws:*:*:%s:*", awsAcctID)
	sourceARNInvalid := "arn:aws:*:*:111111111111:*"
	tests := []struct {
		policy   *policy.PolicyStatement
		expected bool
	}{
		{
			&policy.PolicyStatement{},
			false,
		},
		{
			nil,
			false,
		},
		{
			&policy.PolicyStatement{
				SID:      "test",
				Effect:   "Allow",
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Deny",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"sqs:SendMessage"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"arn:aws:iam::123456:root"},
				},
				Action:   []string{"sqs:SendMessage", "sqs:ReceiveMessage", "sqs:GetQueueAttributes"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"arn:aws:iam::123456789012:root"},
				},
				Action:   []string{"sqs:ReceiveMessage", "sqs:GetQueueAttributes"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"sqs:ReceiveMessage", "sqs:GetQueueAttributes"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"*"},
				Resource: "arn:aws:sqs:us-east-2:444455556666:eventbux-test",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action: []string{"sqs:SendMessage"},
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action: []string{"sqs:SendMessage"},
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNInvalid},
					},
				},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action: []string{"sqs:SendMessage"},
			},
			false,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action: []string{"sqs:SendMessage"},
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{},
				},
			},
			false,
		},

		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"sqs:SendMessage", "sqs:GetQueueAttributes"},
				Resource: "arn:aws:sqs:us-east-2:444455556666:eventbus-test",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			true,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"sqs:SendMessage", "sqs:ReceiveMessage", "sqs:GetQueueAttributes"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			true,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"sqs:SendMessage"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			true,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"SQS:SendMessage"},
				Resource: "*",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			true,
		},
		{
			&policy.PolicyStatement{
				SID:    "test",
				Effect: "Allow",
				Principal: &policy.PolicyStatementPrincipal{
					AWS: []string{"*"},
				},
				Action:   []string{"*"},
				Resource: "arn:aws:sqs:us-east-2:444455556666:eventbus-test",
				Condition: &policy.PolicyStatementCondition{
					ArnEquals: map[string]*policy.OneOrManyString{
						"aws:SourceArn": {sourceARNValid},
					},
				},
			},
			true,
		},
	}

	for _, test := range tests {
		actual := containsEventBusPublishPerms(test.policy, awsAcctID, queueARN)
		if test.expected {
			assert.True(t, actual, "expected policy to succeed")
		} else {
			assert.False(t, actual, "expected policy to fail")
		}
	}
}

func TestValidateQueue(t *testing.T) {
	ctx := context.Background()
	awsAcctID := "123456789012"
	queueURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-asdf-asdf-asdf"
	masterKMSKey := "arn:aws:kms:us-west-2:123456789012:key/zxcv-zxcv-xzcv-xzcv"
	queueARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-asdf-asdf-asdf"

	mockSQSManager := &mocks.SQSManager{}
	qv := &QueueValidator{
		EventBusAWSAccountID:   awsAcctID,
		EncryptionAtRestKeyARN: masterKMSKey,
		SQSManager:             mockSQSManager,
	}

	t.Run("ValidQueue", func(t *testing.T) {
		policy := `
{
  "Version": "2008-10-17",
  "Id": "test-test-test",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "AWS": "*"
      },
      "Action": "sqs:SendMessage",
	  "Resource": "arn:aws:sqs:us-west-2:444422221111:eventbus-asdf-asdf-asdf",
	  "Condition":{
		"ArnEquals":{
		  "aws:SourceArn": "arn:aws:*:*:123456789012:*"
		}
	  }
    }
  ]
}`
		ctx := testCtx(ctx, "valid-queue")
		mockSQSManager.On("GetQueueAttributes", ctx, queueURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"Policy":         policy,
			"QueueArn":       queueARN,
		}, nil)

		err := qv.ValidateQueue(ctx, queueURL)
		assert.NoError(t, err)
	})

	t.Run("InaccessibleAttributes", func(t *testing.T) {
		ctx := testCtx(ctx, "inaccessible-attributes")
		mockSQSManager.On("GetQueueAttributes", ctx, queueURL).Return(nil, errors.New("oh fiddlesticks"))

		err := qv.ValidateQueue(ctx, queueURL)
		assert.Equal(t, ErrInvalidQueueAttributesInaccessible, err)
	})

	t.Run("MissingKMSKey", func(t *testing.T) {
		policy := `
	{
	  "Version": "2008-10-17",
	  "Id": "test-test-test",
	  "Statement": [
	    {
	      "Effect": "Allow",
	      "Principal": "*",
	      "Action": "sqs:SendMessage",
		  "Resource": "arn:aws:sqs:us-west-2:444422221111:eventbus-asdf-asdf-asdf",
		  "Condition":{
			"ArnEquals":{
			  "aws:SourceArn": "arn:aws:*:*:123456789012:*"
			}
		  }
	    }
	  ]
	}`
		ctx := testCtx(ctx, "missing-key")
		mockSQSManager.On("GetQueueAttributes", ctx, queueURL).Return(map[string]string{
			"Policy":   policy,
			"QueueArn": queueARN,
		}, nil)

		err := qv.ValidateQueue(ctx, queueURL)
		assert.Equal(t, ErrInvalidQueueMissingKMSKey, err)
	})

	t.Run("InvalidEventbusQueueURL", func(t *testing.T) {
		invalidQueueURL := "https://sqs.us-west-2.amazonaws.com/444422221111/asdf-asdf-asdf"
		err := qv.ValidateQueue(ctx, invalidQueueURL)
		assert.Equal(t, ErrInvalidQueueMissingEventBusInName, err)
	})

	t.Run("InvalidEventbusPolicyPublishPerms", func(t *testing.T) {
		policy := `
	{
	  "Version": "2008-10-17",
	  "Id": "test-test-test",
	  "Statement": [
	    {
	      "Effect": "Allow",
	      "Principal": "*",
	      "Action": "sqs:ReceiveMessage",
		  "Resource": "arn:aws:sqs:us-west-2:444422221111:eventbus-asdf-asdf-asdf",
		  "Condition":{
			"ArnEquals":{
			  "aws:SourceArn": "arn:aws:*:*:123456789012:*"
			}
		  }
	    }
	  ]
	}`
		ctx := testCtx(ctx, "incorrect-policy")
		mockSQSManager.On("GetQueueAttributes", ctx, queueURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"Policy":         policy,
			"QueueArn":       queueARN,
		}, nil)

		err := qv.ValidateQueue(ctx, queueURL)
		assert.Equal(t, ErrInvalidQueueMissingPublishPermissions, err)
	})

	t.Run("InvalidEventbusPolicyCondition", func(t *testing.T) {
		policy := `
	{
	  "Version": "2008-10-17",
	  "Id": "test-test-test",
	  "Statement": [
	    {
	      "Effect": "Allow",
	      "Principal": "*",
	      "Action": "sqs:SendMessage",
		  "Resource": "*"
	    }
	  ]
	}`
		ctx := testCtx(ctx, "missing-policy-condition")
		mockSQSManager.On("GetQueueAttributes", ctx, queueURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"Policy":         policy,
			"QueueArn":       queueARN,
		}, nil)

		err := qv.ValidateQueue(ctx, queueURL)
		assert.Equal(t, ErrInvalidQueueMissingPublishPermissions, err)
	})
}

func TestValidateDeadletterQueue(t *testing.T) {
	ctx := context.Background()
	dlqURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-deadletter"
	masterKMSKey := "arn:aws:kms:us-west-2:123456789012:key/zxcv-zxcv-xzcv-xzcv"
	dlqARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-deadletter"
	dlqName := "eventbus-deadletter"
	dlqAccountID := "444422221111"
	invalidDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbux-deadletter"
	awsAcctID := "123456789012"

	mockSQSManager := &mocks.SQSManager{}
	qv := &QueueValidator{
		EventBusAWSAccountID:   awsAcctID,
		EncryptionAtRestKeyARN: masterKMSKey,
		SQSManager:             mockSQSManager,
	}

	t.Run("ValidDeadletterQueue", func(t *testing.T) {
		ctx := testCtx(ctx, "valid-dlq")

		mockSQSManager.On("GetQueueURL", ctx, dlqName, dlqAccountID).Return(dlqURL, nil)
		mockSQSManager.On("GetQueueAttributes", ctx, dlqURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"QueueArn":       dlqARN,
		}, nil)

		err := qv.ValidateDeadletterQueue(ctx, dlqARN)
		assert.NoError(t, err)
	})

	t.Run("InvalidDeadletterQueueInaccessibleAttributes", func(t *testing.T) {
		ctx := testCtx(ctx, "dlq-attrs-inaccessible")

		mockSQSManager.On("GetQueueURL", ctx, dlqName, dlqAccountID).Return(dlqURL, nil)
		mockSQSManager.On("GetQueueAttributes", ctx, dlqURL).Return(nil, errors.New("whoa whoa whoa"))

		err := qv.ValidateDeadletterQueue(ctx, dlqARN)
		assert.Equal(t, ErrInvalidDeadletterQueueAttributesInaccessible, err)
	})

	t.Run("InvalidDeadletterQueueMissingKMSMasterKey", func(t *testing.T) {
		ctx := testCtx(ctx, "dlq-missing-kms")

		mockSQSManager.On("GetQueueURL", ctx, dlqName, dlqAccountID).Return(dlqURL, nil)
		mockSQSManager.On("GetQueueAttributes", ctx, dlqURL).Return(map[string]string{
			"QueueArn": dlqARN,
		}, nil)

		err := qv.ValidateDeadletterQueue(ctx, dlqARN)
		assert.Equal(t, ErrInvalidDeadletterQueueKMSKey, err)
	})

	t.Run("ValidDeadletterChainedQueues", func(t *testing.T) {
		ctx := testCtx(ctx, "valid-chained-dlqs")

		accountID := "444422221111"
		firstDLQURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-queue-1-dlq"
		firstDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-queue-1-dlq"
		firstDLQName := "eventbus-queue-1-dlq"

		secondDLQURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-queue-2-dlq"
		secondDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-queue-2-dlq"
		secondDLQName := "eventbus-queue-2-dlq"
		redrive := fmt.Sprintf("{\"deadLetterTargetArn\":\"%s\",\"maxReceiveCount\":5}", secondDLQARN)

		mockSQSManager.On("GetQueueURL", ctx, firstDLQName, accountID).Return(firstDLQURL, nil)
		mockSQSManager.On("GetQueueURL", ctx, secondDLQName, accountID).Return(secondDLQURL, nil)

		mockSQSManager.On("GetQueueAttributes", ctx, firstDLQURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"QueueArn":       firstDLQARN,
			"RedrivePolicy":  redrive,
		}, nil)

		mockSQSManager.On("GetQueueAttributes", ctx, secondDLQURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"QueueArn":       secondDLQARN,
		}, nil)

		err := qv.ValidateDeadletterQueue(ctx, firstDLQARN)
		assert.NoError(t, err)
	})

	t.Run("InvalidDeadletterChainedQueuesMissingKMSMasterKeyOnSecondQueue", func(t *testing.T) {
		ctx := testCtx(ctx, "dlq-chained-missing-kms")

		accountID := "444422221111"
		firstDLQURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-queue-1-dlq"
		firstDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-queue-1-dlq"
		firstDLQName := "eventbus-queue-1-dlq"
		secondDLQURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-queue-2-dlq"
		secondDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-queue-2-dlq"
		secondDLQName := "eventbus-queue-2-dlq"
		redrive := fmt.Sprintf("{\"deadLetterTargetArn\":\"%s\",\"maxReceiveCount\":5}", secondDLQARN)

		mockSQSManager.On("GetQueueURL", ctx, firstDLQName, accountID).Return(firstDLQURL, nil)
		mockSQSManager.On("GetQueueURL", ctx, secondDLQName, accountID).Return(secondDLQURL, nil)

		mockSQSManager.On("GetQueueAttributes", ctx, firstDLQURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"QueueArn":       firstDLQARN,
			"RedrivePolicy":  redrive,
		}, nil)
		mockSQSManager.On("GetQueueAttributes", ctx, secondDLQURL).Return(map[string]string{
			"QueueArn":      firstDLQARN,
			"RedrivePolicy": redrive,
		}, nil)

		err := qv.ValidateDeadletterQueue(ctx, firstDLQARN)
		assert.Equal(t, ErrInvalidDeadletterQueueKMSKey, err)
	})

	t.Run("InvalidDeadletterQueueMissingEventBusInName", func(t *testing.T) {
		err := qv.ValidateDeadletterQueue(ctx, invalidDLQARN)
		assert.Equal(t, ErrInvalidDeadletterQueueMissingEventBusInName, err)
	})

	t.Run("InvalidChainedDeadletterQueueMissingEventBusInName", func(t *testing.T) {
		ctx := testCtx(ctx, "dlq-chained-bad-name")

		accountID := "444422221111"
		firstDLQURL := "https://sqs.us-west-2.amazonaws.com/444422221111/eventbus-queue-1-dlq"
		firstDLQARN := "arn:aws:sqs:us-west-2:444422221111:eventbus-queue-1-dlq"
		firstDLQName := "eventbus-queue-1-dlq"
		secondDLQARN := "arn:aws:sqs:us-west-2:444422221111:queue-2-dlq"
		redrive := fmt.Sprintf("{\"deadLetterTargetArn\":\"%s\",\"maxReceiveCount\":5}", secondDLQARN)

		mockSQSManager.On("GetQueueURL", ctx, firstDLQName, accountID).Return(firstDLQURL, nil)

		mockSQSManager.On("GetQueueAttributes", ctx, firstDLQURL).Return(map[string]string{
			"KmsMasterKeyId": masterKMSKey,
			"QueueArn":       firstDLQARN,
			"RedrivePolicy":  redrive,
		}, nil)

		err := qv.ValidateDeadletterQueue(ctx, firstDLQARN)
		assert.Equal(t, ErrInvalidDeadletterQueueMissingEventBusInName, err)
	})
}

func TestValidQueueName(t *testing.T) {
	tests := []struct {
		name     string
		expected bool
	}{
		{"eventbus", true},
		{"EventBus", true},

		{"eventBus", false},
		{"EVENTBUS", false},
		{"event-bus", false},
		{"event_bus", false},
	}

	for _, test := range tests {
		valid := validQueueName(test.name)
		assert.Equal(t, test.expected, valid)
	}
}

func TestQueueNameFromURL(t *testing.T) {
	tests := []struct {
		url  string
		name string
	}{
		{"https://sqs.us-west-2.amazonaws.com/000123456789/", ""},
		{"https://sqs.us-east-1.amazonaws.com/133713371337/foo", "foo"},
		{"https://sqs.eu-east-1.amazonaws.com/900090009000/bar", "bar"},
	}

	for _, test := range tests {
		name, err := queueNameFromURL(test.url)
		require.NoError(t, err)
		assert.Equal(t, test.name, name)
	}
}

func TestQueueNameFromARN(t *testing.T) {
	tests := []struct {
		arn  string
		name string
	}{
		{"arn:aws:sqs:us-west-2:000123456789:", ""},
		{"arn:aws:sqs:us-east-1:133713371337:foo", "foo"},
		{"arn:aws:sqs:eu-west-1:900090009000:bar", "bar"},
	}

	for _, test := range tests {
		name, err := queueNameFromARN(test.arn)
		require.NoError(t, err)
		assert.Equal(t, test.name, name)
	}
}
