package security

import (
	"database/sql"
	"fmt"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/workload"
	. "github.com/pingcap/endless/testcase/geopartition/utils"
)

func Test(ctx *SuiteContext) {
	It("security test", func() {
		clean := func() {
			MustExec(ctx.DB, "DROP DATABASE IF EXISTS security")
			MustExec(ctx.DB, "DROP DATABASE IF EXISTS security_1")
			MustExec(ctx.DB, "DROP DATABASE IF EXISTS security_2")
			MustExec(ctx.DB, "DROP USER IF EXISTS normal_user")
			MustExec(ctx.DB, "DROP USER IF EXISTS placement_user")
			for i := 0; i < 5; i++ {
				MustExec(ctx.DB, fmt.Sprintf("DROP PLACEMENT POLICY IF EXISTS p%d", i))
			}
		}
		defer clean()
		clean()

		MustExec(ctx.DB, "CREATE DATABASE security")
		MustExec(ctx.DB, "CREATE USER normal_user")
		MustExec(ctx.DB, "CREATE USER placement_user")

		MustExec(ctx.DB, "GRANT SELECT, DROP ON security.* TO normal_user")
		MustExec(ctx.DB, "GRANT SELECT ON security.* TO placement_user")
		MustExec(ctx.DB, "GRANT PLACEMENT_ADMIN ON *.* TO placement_user")

		// security 1: policy manipulation

		// - no priv
		ndb := Try(GetDBWithUser(ctx.TC, "security", "normal_user")).(*sql.DB)
		MustExecError(ndb, 1227, "CREATE PLACEMENT POLICY p1 FOLLOWERS=4")
		MustExecError(ndb, 1227, "ALTER PLACEMENT POLICY IF EXISTS p1 FOLLOWERS=2")
		MustExecError(ndb, 1227, "DROP PLACEMENT POLICY p1")

		// - placement admin
		pdb := Try(GetDBWithUser(ctx.TC, "security", "placement_user")).(*sql.DB)
		MustExec(pdb, "CREATE PLACEMENT POLICY p1 FOLLOWERS=4")
		MustExec(pdb, "ALTER PLACEMENT POLICY IF EXISTS p1 FOLLOWERS=2")
		MustExec(pdb, "DROP PLACEMENT POLICY p1")

		// - super
		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p1 FOLLOWERS=1")
		MustExec(ctx.DB, "CREATE PLACEMENT POLICY p2 FOLLOWERS=8")

		// security 2: binding policy
		MustExec(ctx.DB, "GRANT CREATE, ALTER ON security.* TO normal_user")

		// - table
		MustExec(ndb, "CREATE TABLE t1 (a int) PLACEMENT POLICY = p1")
		MustExec(ndb, "ALTER TABLE t1 PLACEMENT POLICY = p2")
		MustExecError(pdb, 1142, "CREATE TABLE t2 (a int) PLACEMENT POLICY = p1")
		MustExecError(pdb, 1146, "ALTER TABLE t2 PLACEMENT POLICY = p2")

		MustExec(ndb, "DROP TABLE IF EXISTS t1, t2")

		// - partition
		MustExec(ndb, "set @@session.tidb_enable_list_partition = ON")
		MustExec(ndb, `CREATE TABLE t1 (id INT)
PLACEMENT POLICY = p1
PARTITION BY RANGE (id) (
    PARTITION p0 VALUES LESS THAN (6),
    PARTITION p1 VALUES LESS THAN (16) PLACEMENT POLICY = p2
)`)
		MustExec(ndb, `CREATE TABLE t2 (id INT) PLACEMENT POLICY = p2 PARTITION BY HASH(id)`)
		MustExec(ndb, `CREATE TABLE t3 (id int, name varchar(10))
PLACEMENT POLICY = p1
PARTITION BY LIST COLUMNS(id,name) (
     partition p0 values IN ((1,'a'),(2,'b')) PLACEMENT POLICY = p2,
     partition p1 values IN ((3,'c'),(4,'d')),
     partition p3 values IN ((5,'e'),(null,null))
)`)
		MustExec(ndb, `ALTER TABLE t1
ADD PARTITION (
    PARTITION p5 VALUES LESS THAN (36) PLACEMENT POLICY = p1
)`)
		MustExec(ndb, `ALTER TABLE t3
ADD PARTITION (
     partition p4 values IN ((9,'a'),(10,'b')) PLACEMENT POLICY = p1
)`)

		MustExec(pdb, "set @@session.tidb_enable_list_partition = ON")
		MustExecError(pdb, 1142, `CREATE TABLE t4 (id INT)
PLACEMENT POLICY = p1
PARTITION BY RANGE (id) (
    PARTITION p0 VALUES LESS THAN (6),
    PARTITION p1 VALUES LESS THAN (16) PLACEMENT POLICY = p2
)`)
		MustExecError(pdb, 1142, `CREATE TABLE t5 (id INT) PLACEMENT POLICY = p2 PARTITION BY HASH(id)`)
		MustExecError(pdb, 1142, `CREATE TABLE t6 (id int, name varchar(10))
PLACEMENT POLICY = p1
PARTITION BY LIST COLUMNS(id,name) (
     partition p0 values IN ((1,'a'),(2,'b')) PLACEMENT POLICY = p2,
     partition p1 values IN ((3,'c'),(4,'d')),
     partition p3 values IN ((5,'e'),(null,null))
)`)
		MustExecError(pdb, 1146, `ALTER TABLE t4
ADD PARTITION (
    PARTITION p5 VALUES LESS THAN (36) PLACEMENT POLICY = p1
)`)
		MustExecError(pdb, 1146, `ALTER TABLE t6
ADD PARTITION (
     partition p4 values IN ((9,'a'),(10,'b')) PLACEMENT POLICY = p1
)`)
		MustExec(ndb, "DROP TABLE IF EXISTS t1,t2,t3,t4,t5,t6")

		// - database
		MustExec(ctx.DB, "GRANT CREATE, ALTER ON security_1.* TO normal_user")
		MustExec(ndb, "CREATE DATABASE security_1 PLACEMENT POLICY = p1")
		MustExec(ndb, "ALTER DATABASE security_1 PLACEMENT POLICY = p2")
		MustExecError(pdb, 1044, "CREATE DATABASE security_2 PLACEMENT POLICY = p1")
		MustExecError(pdb, 1044, "ALTER DATABASE security_2 PLACEMENT POLICY = p2")
		MustExec(ctx.DB, "REVOKE CREATE, ALTER ON security_1.* FROM normal_user")

		// security 3: show placement
		MustExec(ndb, "CREATE TABLE t1 (a int) PLACEMENT POLICY = p1")
		MustExec(ndb, "CREATE TABLE t2 (a int) PLACEMENT POLICY = p2")
		MustExec(ndb, `CREATE TABLE t3 (id INT)
PARTITION BY RANGE (id) (
    PARTITION p0 VALUES LESS THAN (6),
    PARTITION p1 VALUES LESS THAN (16) PLACEMENT POLICY = p2
)`)
		Try(WaitScheduledOrPending(ndb, "t1", time.Minute))
		Try(WaitScheduledOrPending(ndb, "t2", time.Minute))
		Try(WaitScheduledOrPending(ndb, "t3", time.Minute))
		time.Sleep(time.Second * 2)
		MustExec(ctx.DB, "REVOKE ALL PRIVILEGES ON security.* FROM normal_user")
		MustExec(ctx.DB, "FLUSH PRIVILEGES")

		privs := []string{
			"create",
			"alter",
			"drop",
			"select",
			"insert",
			"references",
			"index",
			"all privileges",
		}
		// for one table
		for _, priv := range privs {
			MustExec(ctx.DB, fmt.Sprintf("GRANT %s ON security.t1 TO normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
				{"TABLE security.t1", "FOLLOWERS=1", "SCHEDULED"},
			}))
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT FOR TABLE t1"))).To(Equal([][]string{
				{"TABLE security.t1", "FOLLOWERS=1", "SCHEDULED"},
			}))
			MustExec(ctx.DB, fmt.Sprintf("REVOKE %s ON security.t1 FROM normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
			}))
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t1")
		}
		// for one database
		for _, priv := range privs {
			MustExec(ctx.DB, fmt.Sprintf("GRANT %s ON security.* TO normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
				{"TABLE security.t1", "FOLLOWERS=1", "SCHEDULED"},
				{"TABLE security.t2", "FOLLOWERS=8", "PENDING"},
				{"TABLE security.t3 PARTITION p1", "FOLLOWERS=8", "PENDING"},
			}))
			MustExec(ctx.DB, fmt.Sprintf("REVOKE %s ON security.* FROM normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
			}))
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t1")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t2")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t3")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t3 PARTITION p1")
		}
		// for all
		for _, priv := range privs {
			MustExec(ctx.DB, fmt.Sprintf("GRANT %s ON *.* TO normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
				{"DATABASE security_1", "FOLLOWERS=8", "SCHEDULED"},
				{"TABLE security.t1", "FOLLOWERS=1", "SCHEDULED"},
				{"TABLE security.t2", "FOLLOWERS=8", "PENDING"},
				{"TABLE security.t3 PARTITION p1", "FOLLOWERS=8", "PENDING"},
			}))
			MustExec(ctx.DB, fmt.Sprintf("REVOKE %s ON *.* FROM normal_user", priv))
			MustExec(ctx.DB, "FLUSH PRIVILEGES")
			Expect(ReadAll(MustQuery(ndb, "SHOW PLACEMENT"))).To(Equal([][]string{
				{"POLICY p1", "FOLLOWERS=1", "NULL"},
				{"POLICY p2", "FOLLOWERS=8", "NULL"},
			}))
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t1")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t2")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t3")
			MustExecError(ndb, 1142, "SHOW PLACEMENT FOR TABLE t3 PARTITION p1")
			MustExecError(ndb, 1064, "SHOW PLACEMENT FOR security_1")
		}

		MustExec(ctx.DB, "DROP TABLE IF EXISTS t1,t2")

		// security 4: show create placement policy
		MustExec(ctx.DB, "REVOKE ALL PRIVILEGES ON *.* FROM normal_user")
		MustExec(ctx.DB, "FLUSH PRIVILEGES")
		Expect(ReadAll(MustQuery(ndb, "SHOW CREATE PLACEMENT POLICY p1"))).To(Equal([][]string{
			{"p1", "CREATE PLACEMENT POLICY `p1` FOLLOWERS=1"},
		}))
		Expect(ReadAll(MustQuery(ndb, "SHOW CREATE PLACEMENT POLICY p2"))).To(Equal([][]string{
			{"p2", "CREATE PLACEMENT POLICY `p2` FOLLOWERS=8"},
		}))
		MustExec(pdb, "SHOW CREATE PLACEMENT POLICY p1")
		Expect(ReadAll(MustQuery(pdb, "SHOW CREATE PLACEMENT POLICY p1"))).To(Equal([][]string{
			{"p1", "CREATE PLACEMENT POLICY `p1` FOLLOWERS=1"},
		}))
		MustExec(pdb, "SHOW CREATE PLACEMENT POLICY p2")
		Expect(ReadAll(MustQuery(pdb, "SHOW CREATE PLACEMENT POLICY p2"))).To(Equal([][]string{
			{"p2", "CREATE PLACEMENT POLICY `p2` FOLLOWERS=8"},
		}))
		MustExec(ctx.DB, "SHOW CREATE PLACEMENT POLICY p1")
		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE PLACEMENT POLICY p1"))).To(Equal([][]string{
			{"p1", "CREATE PLACEMENT POLICY `p1` FOLLOWERS=1"},
		}))
		MustExec(ctx.DB, "SHOW CREATE PLACEMENT POLICY p2")
		Expect(ReadAll(MustQuery(ctx.DB, "SHOW CREATE PLACEMENT POLICY p2"))).To(Equal([][]string{
			{"p2", "CREATE PLACEMENT POLICY `p2` FOLLOWERS=8"},
		}))

		// security 5: information.placement_rules
		data := Sort(ReadAll(MustQuery(ctx.DB, "SELECT * FROM INFORMATION_SCHEMA.PLACEMENT_POLICIES")))
		Expect(data[0][2]).To(Equal("p1"))
		Expect(data[1][2]).To(Equal("p2"))

		data = Sort(ReadAll(MustQuery(ndb, "SELECT * FROM INFORMATION_SCHEMA.PLACEMENT_POLICIES")))
		Expect(data[0][2]).To(Equal("p1"))
		Expect(data[1][2]).To(Equal("p2"))

		data = Sort(ReadAll(MustQuery(pdb, "SELECT * FROM INFORMATION_SCHEMA.PLACEMENT_POLICIES")))
		Expect(data[0][2]).To(Equal("p1"))
		Expect(data[1][2]).To(Equal("p2"))
	})
}