package main

import (
	"context"
	"sync"

	"github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/plugin"
	"github.com/pingcap/tidb/pkg/sessionctx/variable"
	"go.uber.org/zap"
)

func init() {
	variable.RegisterSysVar(redactLogSv)
	variable.RegisterSysVar(auditLogMaxSizeSv)
	variable.RegisterSysVar(auditLogGeneralEvents)
	variable.RegisterSysVar(auditLogIncludeExtrasSv)
}

// OnInit implements TiDB plugin's OnInit SPI.
func OnInit(ctx context.Context, manifest *plugin.Manifest) error {
	err := global.init()
	if err != nil {
		return err
	}
	return initLog()
}

// OnFlush implements TiDB plugin's OnFlush SPI.
func OnFlush(ctx context.Context, manifest *plugin.Manifest) error {
	return global.update()
}

// OnShutdown implements TiDB plugin's OnShutdown SPI.
func OnShutdown(ctx context.Context, manifest *plugin.Manifest) error {
	log.Info("Sync audit log when TiDB shutdown")
	logger.Load().(*zap.Logger).Sync()
	return nil
}

// OnConnectionEvent implements TiDB Audit plugin's OnConnectionEvent SPI.
func OnConnectionEvent(ctx context.Context, event plugin.ConnectionEvent, connInfo *variable.ConnectionInfo) error {
	if event == plugin.PreAuth {
		return nil
	}
	var reason string
	if r := ctx.Value(plugin.RejectReasonCtxValue{}); r != nil {
		reason = r.(string)
	}
	if !global.get().userIsAuditable(connInfo.User) && reason == "" {
		// The user is not auditable (i.e. its an application connection)
		// We skip auditing provided there is no rejection reason.
		return nil
	}
	audit(ctx, "CONNECTION", event.String(), 0).
		fields([]zap.Field{
			zap.String("HOST", connInfo.Host),
			zap.String("CLIENT_IP", connInfo.ClientIP),
			zap.String("USER", connInfo.User),
			zap.Strings("DATABASES", []string{connInfo.DB}),
			zap.Strings("TABLES", []string{}),
			zap.String("SQL_TEXT", ""),
			zap.Uint64("ROWS", 0),
			zap.String("CLIENT_PORT", connInfo.ClientPort),
			zap.Uint64("CONNECTION_ID", uint64(connInfo.ConnectionID)),
			zap.String("CONNECTION_TYPE", connInfo.ConnectionType),
			zap.Int("SERVER_ID", connInfo.ServerID),
			zap.Int("SERVER_PORT", connInfo.ServerPort),
			zap.Float64("DURATION", connInfo.Duration),
			zap.String("SERVER_OS_LOGIN_USER", connInfo.ServerOSLoginUser),
			zap.String("OS_VERSION", connInfo.OSVersion),
			zap.String("CLIENT_VERSION", connInfo.ClientVersion),
			zap.String("SERVER_VERSION", connInfo.ServerVersion),
			zap.String("AUDIT_VERSION", ""),
			zap.String("SSL_VERSION", connInfo.SSLVersion),
			zap.Int("PID", connInfo.PID),
			zap.String("Reason", reason),
		}).
		log()
	return nil
}

// OnGeneralEvent implements TiDB Audit plugin's OnGeneralEvent SPI.
func OnGeneralEvent(ctx context.Context, sctx *variable.SessionVars, event plugin.GeneralEvent, cmd string) {
	switch cmd {
	case "Close stmt":
		if !auditLogExtrasIncludes(auditLogExtrasClose) {
			return
		}
	}

	connInfo := sctx.ConnectionInfo
	if connInfo == nil {
		return
	}

	if !auditLogEnabledEvents.IsGeneralEventEnabled(event) {
		return
	}

	// Check if the user is not auditable
	// i.e. it is common to not audit application users.
	if !global.get().userIsAuditable(connInfo.User) {
		return
	}
	setToSlice := func(set map[string]struct{}) (slice []string) {
		slice = make([]string, 0, len(set))
		for val := range set {
			slice = append(slice, val)
		}
		return
	}
	var (
		tableNames []string
		dbNames    = make(map[string]struct{})
	)
	eventClass, eventSubClass := "GENERAL", ""
	if len(sctx.StmtCtx.Tables) > 0 {
		eventClass = "TABLE_ACCESS"
		eventSubClass = sctx.StmtCtx.StmtType
		for _, tbl := range sctx.StmtCtx.Tables {
			dbNames[tbl.DB] = struct{}{}
			tableNames = append(tableNames, tbl.Table)
			if !global.get().check(connInfo.User, tbl.DB, tbl.Table, eventSubClass) {
				return
			}
		}
	}
	stmt := sctx.StmtCtx.OriginalSQL
	if normalizeStatements() || isPasswordStmt(sctx) {
		stmt, _ = sctx.StmtCtx.SQLDigest()
	}

	var paramStrs []string
	if !normalizeStatements() && sctx.PlanCacheParams != nil {
		params := sctx.PlanCacheParams.AllParamValues()
		paramStrs = make([]string, 0, len(params))
		for _, param := range params {
			// param is formated with type info, e.g. `KindInt64 12`
			paramStrs = append(paramStrs, param.String())
		}
	}

	fields := []zap.Field{
		zap.String("HOST", connInfo.Host),
		zap.String("CLIENT_IP", connInfo.ClientIP),
		zap.String("USER", connInfo.User),
		zap.Strings("DATABASES", setToSlice(dbNames)),
		zap.Strings("TABLES", tableNames),
		zap.String("SQL_TEXT", stmt),
		zap.Uint64("ROWS", sctx.StmtCtx.AffectedRows()),
		zap.Uint64("CONNECTION_ID", uint64(connInfo.ConnectionID)),
		zap.String("CLIENT_PORT", connInfo.ClientPort),
		zap.Int("PID", connInfo.PID),
		zap.String("COMMAND", cmd),
		zap.String("SQL_STATEMENTS", sctx.StmtCtx.StmtType),
		zap.Strings("EXECUTE_PARAMS", paramStrs),
		// Don't use `connInfo.DB`, as it represents the default database of the connection, but
		// not the current database of this session.
		zap.String("CURRENT_DB", sctx.CurrentDB),
		zap.String("EVENT", event.String()),
	}
	if id := ctx.Value(plugin.PrepareStmtIDCtxKey); id != nil {
		fields = append(fields, zap.Uint32("PREPARED_STMT_ID", id.(uint32)))
	}
	if retrying := ctx.Value(plugin.IsRetryingCtxKey); retrying != nil && retrying.(bool) {
		// Only print the RETRY field when it's true
		fields = append(fields, zap.Bool("RETRY", true))
	}

	audit(ctx, eventClass, eventSubClass, 0).
		fields(fields).
		log()
	return
}

var parserPool = &sync.Pool{New: func() interface{} { return parser.New() }}

func isPasswordStmt(vars *variable.SessionVars) bool {
	if vars.StmtCtx.OriginalSQL == "" {
		return false
	}

	var stmts []ast.StmtNode
	switch vars.StmtCtx.StmtType {
	case "CreateUser", "Set", "other", "":
	default:
		return false
	}

	p := parserPool.Get().(*parser.Parser)
	defer parserPool.Put(p)
	stmts, _, _ = p.ParseSQL(vars.StmtCtx.OriginalSQL)

	findInSpecs := func(specs []*ast.UserSpec) bool {
		for _, spec := range specs {
			if authOpt := spec.AuthOpt; authOpt != nil && (authOpt.ByAuthString || authOpt.ByHashString) {
				return true
			}
		}
		return false
	}

	for _, node := range stmts {
		switch stmt := node.(type) {
		case *ast.SetPwdStmt:
			return true
		case *ast.CreateUserStmt:
			if findInSpecs(stmt.Specs) {
				return true
			}
		case *ast.AlterUserStmt:
			if findInSpecs(stmt.Specs) {
				return true
			}
		}
	}
	return false
}