package audit

import (
	"context"
	"fmt"
	"path"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/pingcap/tidb/pkg/errno"
	"github.com/pingcap/tidb/pkg/parser/auth"
	"github.com/pingcap/tidb/pkg/server"
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/pingcap/tidb/pkg/util/sem"
	"github.com/stretchr/testify/require"
)

func eventClasses2String(eventClasses []EventClass) string {
	eventClassNames := make([]string, len(eventClasses))
	for i, eventClass := range eventClasses {
		eventClassNames[i] = class2String[eventClass]
	}
	return fmt.Sprintf(`[EVENT="[%s]"]`, strings.Join(eventClassNames, ","))
}

func TestAuditLogRotate(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)

	tempDir := t.TempDir()
	tk.MustExec("SET global tidb_audit_log = ?", path.Join(tempDir, DefAuditLogName))
	_, err := deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)

	tk.MustExec("SET global tidb_audit_log_reserved_backups = 2")
	tk.MustExec("SET global tidb_audit_enabled = 1")
	for i := 0; i < 5; i++ {
		tk.MustQuery("SELECT audit_log_rotate()").Check(testkit.Rows("OK"))
		time.Sleep(time.Second)
	}

	files, err := deleteAllAuditLogs(tempDir, "tidb-audit-", ".log")
	require.NoError(t, err)
	require.Equal(t, 2, len(files), files)
}

func TestAuditLogEnableRule(t *testing.T) {
	register4Test(t)
	tempDir := t.TempDir()
	store := testkit.CreateMockStore(t)
	srv := server.CreateMockServer(t, store)
	defer srv.Close()
	conn := server.CreateMockConn(t, srv)
	defer conn.Close()
	tk := testkit.NewTestKit(t, store)
	err := tk.Session().Auth(&auth.UserIdentity{Username: "root", Hostname: "%"}, nil, nil, nil)
	require.NoError(t, err)
	tk.MustExec(fmt.Sprintf("SET global tidb_audit_log = '%s'", filepath.Join(tempDir, DefAuditLogName)))
	tk.MustExec("SET global tidb_audit_enabled = 1")

	_, err = deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)

	tk.MustExec("use test")
	tk.MustExec("drop table if exists t")
	tk.MustExec("create table t (c int)")
	tk.MustQuery("select audit_log_create_filter('all', '{}')").Check(testkit.Rows("OK"))
	tk.MustQuery("select content from mysql.audit_log_filters where filter_name = 'all'").Check(testkit.Rows(`{}`))
	tk.MustQuery("select audit_log_create_rule('%@%', 'all')").Check(testkit.Rows("OK"))
	tk.MustQuery("select enabled from mysql.audit_log_filter_rules where user = '%@%' and filter_name = 'all'").Check(testkit.Rows("1"))
	tk.MustExec("SET global tidb_audit_log_redacted = 0")

	require.NoError(t, conn.HandleQuery(context.Background(), "use test"))
	require.NoError(t, conn.HandleQuery(context.Background(), "insert into t values (1)"))
	ok, _, err := containsMessage(filepath.Join(tempDir, DefAuditLogName), "insert into t values (1)")
	require.NoError(t, err)
	require.True(t, ok)

	tk.MustQuery("select audit_log_disable_rule('%@%', 'all')").Check(testkit.Rows("OK"))
	require.NoError(t, conn.HandleQuery(context.Background(), "insert into t values (2)"))
	ok, _, err = containsMessage(filepath.Join(tempDir, DefAuditLogName), "insert into t values (2)")
	require.NoError(t, err)
	require.False(t, ok)

	tk.MustQuery("select audit_log_enable_rule('%@%', 'all')").Check(testkit.Rows("OK"))
	require.NoError(t, conn.HandleQuery(context.Background(), "insert into t values (3)"))
	ok, _, err = containsMessage(filepath.Join(tempDir, DefAuditLogName), "insert into t values (3)")
	require.NoError(t, err)
	require.True(t, ok)

	_, err = deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)
}

func TestAuditAdmin(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	srv := server.CreateMockServer(t, store)
	defer srv.Close()
	conn := server.CreateMockConn(t, srv)
	defer conn.Close()
	tk := testkit.NewTestKit(t, store)
	err := tk.Session().Auth(&auth.UserIdentity{Username: "root", Hostname: "%"}, nil, nil, nil)
	require.NoError(t, err)

	tempDir := t.TempDir()
	tk.MustExec("set global tidb_audit_log=?", path.Join(tempDir, DefAuditLogName))
	tk.MustExec("CREATE USER testuser")
	tk2 := testkit.NewTestKit(t, store)
	err = tk2.Session().Auth(&auth.UserIdentity{Username: "testuser", Hostname: "%"}, nil, nil, nil)
	require.NoError(t, err)
	tk.MustExec("GRANT SELECT, SYSTEM_VARIABLES_ADMIN ON *.* to testuser")

	tk2.MustGetErrCode(`select audit_log_create_filter("empty", '{"filter":[]}')`, errno.ErrSpecificAccessDenied)
	tk2.MustGetErrCode(`set global tidb_audit_enabled = 0`, errno.ErrSpecificAccessDenied)
	tk2.MustGetErrCode(`select * from mysql.audit_log_filters`, errno.ErrTableaccessDenied)

	tk.MustExec("GRANT AUDIT_ADMIN ON *.* to testuser")

	tk2.MustQuery(`select audit_log_create_filter("all_query", '{"filter":[{"class":["QUERY"]}]}')`).Check(testkit.Rows("OK"))
	tk2.MustExec(`set global tidb_audit_enabled = 0`)
	tk2.MustQuery(`select * from mysql.audit_log_filters`).Check(testkit.Rows(`all_query {"filter":[{"class":["QUERY"]}]}`))

	_, err = deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)
}

func TestRestrictedAuditAdmin(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	sem.Enable()
	defer sem.Disable()
	srv := server.CreateMockServer(t, store)
	defer srv.Close()
	conn := server.CreateMockConn(t, srv)
	defer conn.Close()
	tk := testkit.NewTestKit(t, store)
	tempDir := t.TempDir()
	tk.MustExec("set global tidb_audit_log=?", path.Join(tempDir, DefAuditLogName))

	tk.MustExec("CREATE USER testuser")
	tk2 := testkit.NewTestKit(t, store)
	err := tk2.Session().Auth(&auth.UserIdentity{Username: "testuser", Hostname: "%"}, nil, nil, nil)
	require.NoError(t, err)
	tk.MustExec("GRANT SELECT, SYSTEM_VARIABLES_ADMIN, AUDIT_ADMIN ON *.* to testuser")

	tk2.MustGetErrCode(`select audit_log_create_filter("empty", '{"filter":[]}')`, errno.ErrSpecificAccessDenied)
	tk2.MustGetErrCode(`set global tidb_audit_enabled = 0`, errno.ErrSpecificAccessDenied)
	tk2.MustGetErrCode(`select * from mysql.audit_log_filters`, errno.ErrTableaccessDenied)

	tk.MustExec("GRANT RESTRICTED_AUDIT_ADMIN ON *.* to testuser")

	tk2.MustQuery(`select audit_log_create_filter("all_query", '{"filter":[{"class":["QUERY"]}]}')`).Check(testkit.Rows("OK"))
	tk2.MustExec(`set global tidb_audit_enabled = 0`)
	tk2.MustQuery(`select * from mysql.audit_log_filters`).Check(testkit.Rows(`all_query {"filter":[{"class":["QUERY"]}]}`))

	_, err = deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)
}

func TestEventClass(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	srv := server.CreateMockServer(t, store)
	defer srv.Close()
	conn := server.CreateMockConn(t, srv)
	defer conn.Close()

	tempDir := t.TempDir()
	logPath := path.Join(tempDir, DefAuditLogName)
	require.NoError(t, conn.HandleQuery(
		context.Background(), fmt.Sprintf("SET global tidb_audit_log='%s'", logPath),
	))

	require.NoError(t, conn.HandleQuery(context.Background(), "SET global tidb_audit_enabled = 1"))
	require.NoError(t, conn.HandleQuery(context.Background(), "SELECT audit_log_create_filter('all', '{}')"))
	require.NoError(t, conn.HandleQuery(context.Background(), "SELECT audit_log_create_rule('%@%', 'all')"))

	// test QUERY and AUDIT
	testcases := []struct {
		sql           string
		requireEvents []EventClass
		success       bool
	}{
		{"use test", []EventClass{ClassQuery}, true},
		{"begin", []EventClass{ClassQuery, ClassTransaction}, true},
		{"create table if not exists t (c int)", []EventClass{ClassQuery, ClassDDL}, true},
		{"select * from t", []EventClass{ClassQuery, ClassSelect}, true},
		{"commit", []EventClass{ClassQuery, ClassTransaction}, true},
		{"prepare abs_func from 'select abs(?)'", []EventClass{ClassQuery}, true},
		{"set @num = -1", []EventClass{ClassQuery}, true},
		{"execute abs_func using @num", []EventClass{ClassQuery, ClassExecute, ClassSelect}, true},
		{"insert into t values (1)", []EventClass{ClassQuery, ClassDML, ClassInsert}, true},
		{"replace into t values (2)", []EventClass{ClassQuery, ClassDML, ClassReplace}, true},
		{"update t set c = 3 where c = 1", []EventClass{ClassQuery, ClassDML, ClassUpdate}, true},
		{"delete from t", []EventClass{ClassQuery, ClassDML, ClassDelete}, true},
		{`LOAD DATA LOCAL INFILE 'load_data.csv' INTO TABLE t FIELDS TERMINATED BY ',' ENCLOSED BY '\"' LINES TERMINATED BY '\r\n'`, []EventClass{ClassQuery, ClassDML, ClassLoadData}, false},
		{"select audit_log_create_filter('empty', '{}')", []EventClass{ClassAudit, ClassAuditFuncCall}, true},
		{"set global tidb_audit_enabled = 1", []EventClass{ClassAudit, ClassAuditSetSysVar, ClassAuditEnable}, true},
		{"set global tidb_audit_enabled = 0", []EventClass{ClassAudit, ClassAuditSetSysVar, ClassAuditDisable}, true},
	}
	for _, tc := range testcases {
		if tc.success {
			require.NoError(t, conn.HandleQuery(context.Background(), tc.sql), tc.sql)
		} else {
			require.Error(t, conn.HandleQuery(context.Background(), tc.sql), tc.sql)
		}
		eventClassesStr := eventClasses2String(tc.requireEvents)
		ok, log, err := containsMessage(logPath, eventClassesStr)
		require.NoError(t, err, "%s\n%s\n%s", tc.sql, eventClassesStr, log)
		require.True(t, ok, "%s\n%s\n%s", tc.sql, eventClassesStr, log)
		require.NoError(t, conn.HandleQuery(context.Background(), "SELECT audit_log_rotate()"), "%s\n%s\n%s", tc.sql, eventClassesStr, log)
		time.Sleep(time.Second)
	}
}

func TestAuditLogRuleUsername(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tempDir := t.TempDir()
	tk.MustExec("set global tidb_audit_log=?", path.Join(tempDir, DefAuditLogName))
	_, err := deleteAllAuditLogs(tempDir, "tidb-audit", ".log")
	require.NoError(t, err)

	tk.MustQuery("select audit_log_create_filter('all', '{}')").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_rule('foo@bar@127.0.0.1', 'all')").Check(testkit.Rows("OK"))
}

func TestAuditLogRuleValidate(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tempDir := t.TempDir()
	tk.MustExec("set global tidb_audit_log=?", path.Join(tempDir, DefAuditLogName))
	tk.MustQuery("select audit_log_create_filter('f1', '{}')").Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters").Check(testkit.Rows("f1 {}"))
	// dup name
	err := tk.QueryToErr("select audit_log_create_filter('f1', '{}')")
	require.EqualError(t, err, "[kv:1062]Duplicate entry 'f1' for key 'audit_log_filters.PRIMARY'")
	// invalid event name
	err = tk.QueryToErr(`select audit_log_create_filter('f2', '{"filter": [{"class": ["noexisteventname"]}]}')`)
	require.EqualError(t, err, "invalid event class name 'noexisteventname'")
	// invalid status code
	err = tk.QueryToErr(`select audit_log_create_filter('f2', '{"filter": [{"status_code": [999]}]}')`)
	require.EqualError(t, err, "invalid status code '999'")
	// invalid filters should not be created
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters").Check(testkit.Rows("f1 {}"))
}

func TestFilterCURD(t *testing.T) {
	register4Test(t)
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tempDir := t.TempDir()
	tk.MustExec("set global tidb_audit_log=?", path.Join(tempDir, DefAuditLogName))

	// filters and rules should be empty for a new cluster
	tk.MustQuery("select * from mysql.audit_log_filters").Check(testkit.Rows())
	tk.MustQuery("select * from mysql.audit_log_filter_rules").Check(testkit.Rows())

	// create filters
	tk.MustQuery("select audit_log_create_filter('f1', '{}')").Check(testkit.Rows("OK"))
	// A not normalized JSON string should be normalized
	tk.MustQuery("select audit_log_create_filter('f2', ?)", `{"filter": [ {"class": [ "query" ]}]}`).
		Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_filter('f3', ?, true)", `{"filter":[{"class":["CONNECTION"]}]}`).
		Check(testkit.Rows("OK"))
	// empty list should be normalized
	tk.MustQuery("select audit_log_create_filter('f4', ?, false)",
		`{"class":[],"table":[],"class_excl":[],"table_excl":[],"status_code":[]}`).
		Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			"f1 {}",
			`f2 {"filter":[{"class":["QUERY"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))

	// should not create dup name filter if replace is false
	err := tk.QueryToErr("select audit_log_create_filter('f2', '{}', false)")
	require.EqualError(t, err, "[kv:1062]Duplicate entry 'f2' for key 'audit_log_filters.PRIMARY'")
	err = tk.QueryToErr("select audit_log_create_filter('f3', '{}')")
	require.EqualError(t, err, "[kv:1062]Duplicate entry 'f3' for key 'audit_log_filters.PRIMARY'")
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			"f1 {}",
			`f2 {"filter":[{"class":["QUERY"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))

	// should replace filter content when replace=true
	tk.MustQuery("select audit_log_create_filter('f1', ?, true)", `{"filter":[{"table":["test.*"]}]}`).
		Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			`f1 {"filter":[{"table":["test.*"]}]}`,
			`f2 {"filter":[{"class":["QUERY"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))

	// create filter rules
	tk.MustQuery("select * from mysql.audit_log_filter_rules").Check(testkit.Rows())
	tk.MustQuery("select audit_log_create_rule('%@%', 'f1')").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_rule('root@%', 'f2', true)").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_rule('user1@localhost', 'f3', false)").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_rule('%@%', 'f4')").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_create_rule('root@%', 'f1')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "root@% f2 1", "user1@localhost f3 1"))

	// disable rule
	tk.MustQuery("select audit_log_disable_rule('%@%', 'f1')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 0", "%@% f4 1", "root@% f1 1", "root@% f2 1", "user1@localhost f3 1"))
	tk.MustQuery("select audit_log_disable_rule('root@%', 'f2')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 0", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))

	// disable again should take no effect
	tk.MustQuery("select audit_log_disable_rule('%@%', 'f1')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 0", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))

	// enable rule
	tk.MustQuery("select audit_log_enable_rule('%@%', 'f1')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))
	tk.MustQuery("select audit_log_enable_rule('user1@localhost', 'f3')").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))

	// enable/disable exists rule should return error
	err = tk.QueryToErr("select audit_log_enable_rule('userx@%', 'f1')")
	require.EqualError(t, err, "rule does not exist")
	err = tk.QueryToErr("select audit_log_disable_rule('%@%', 'f2')")
	require.EqualError(t, err, "rule does not exist")

	// replace rule
	tk.MustQuery("select audit_log_create_rule('%@%', 'f1', true)").Check(testkit.Rows("OK"))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))

	// cannot delete an in-use filter
	err = tk.QueryToErr("select audit_log_remove_filter('f2')")
	require.EqualError(t, err, "filter 'f2' is in use")
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			`f1 {"filter":[{"table":["test.*"]}]}`,
			`f2 {"filter":[{"class":["QUERY"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "root@% f2 0", "user1@localhost f3 1"))

	// delete rule
	tk.MustQuery("select audit_log_remove_rule('root@%', 'f2')").Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			`f1 {"filter":[{"table":["test.*"]}]}`,
			`f2 {"filter":[{"class":["QUERY"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "user1@localhost f3 1"))

	// delete a unused filter
	tk.MustQuery("select audit_log_remove_filter('f2')").Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			`f1 {"filter":[{"table":["test.*"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "user1@localhost f3 1"))

	// delete a no-exist rule or filter should take no effect
	tk.MustQuery("select audit_log_remove_rule('root@%', 'f4')").Check(testkit.Rows("OK"))
	tk.MustQuery("select audit_log_remove_filter('f5')").Check(testkit.Rows("OK"))
	tk.MustQuery("select FILTER_NAME, CONTENT from mysql.audit_log_filters order by FILTER_NAME").
		Check(testkit.Rows(
			`f1 {"filter":[{"table":["test.*"]}]}`,
			`f3 {"filter":[{"class":["CONNECTION"]}]}`,
			`f4 {}`,
		))
	tk.MustQuery("select USER, FILTER_NAME, ENABLED from mysql.audit_log_filter_rules order by USER").
		Check(testkit.Rows("%@% f1 1", "%@% f4 1", "root@% f1 1", "user1@localhost f3 1"))
}