package main

import (
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/tencentcloud/ssm-rotation-sdk-golang/lib/db"
	"github.com/tencentcloud/ssm-rotation-sdk-golang/lib/ssm"
	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
	"log"
	"time"
)

var (
	roleArn, tokenPath, providerId, regionName, saToken string
	secretName, dbAddress, dbName, ssmRegionName        string
	dbPort                                              uint64
	dbConn                                              *db.DynamicSecretRotationDb
	Header                                              = map[string]string{
		"Authorization":      "SKIP",
		"X-TC-Action":        "AssumeRoleWithWebIdentity",
		"Host":               "sts.internal.tencentcloudapi.com",
		"X-TC-RequestClient": "PHP_SDK",
		"X-TC-Version":       "2018-08-13",
		"X-TC-Region":        regionName,
		"X-TC-Timestamp":     "1659944952",
		"Content-type":       "application/json",
	}
)

type Credentials struct {
	TmpSecretId  string
	TmpSecretKey string
	Token        string
	ExpiredTime  uint64
}

func main() {
	flag.StringVar(&secretName, "ssmName", "", "ssm名称")
	flag.StringVar(&ssmRegionName, "ssmRegionName", "", "ssm地域")
	flag.StringVar(&dbAddress, "dbAddress", "", "数据库地址")
	flag.StringVar(&dbName, "dbName", "", "数据库名称")
	flag.Uint64Var(&dbPort, "dbPort", 0, "数据库端口")
	flag.Parse()

	provider, err := common.DefaultTkeOIDCRoleArnProvider()
	if err != nil {
		log.Fatal("failed to assume role with web identity, err:", err)
	}
	assumeResp, err := provider.GetCredential()
	if err != nil {
		log.Fatal("failed to assume role with web identity, err:", err)
	}

	var credential Credentials
	if assumeResp != nil {
		credential = Credentials{
			TmpSecretId:  assumeResp.GetSecretId(),
			TmpSecretKey: assumeResp.GetSecretKey(),
			Token:        assumeResp.GetToken(),
		}
	}
	log.Printf("secretId:%v,secretey%v,token%v\\n", credential.TmpSecretId, credential.TmpSecretKey, credential.Token)
	DB(credential)
}

func DB(credential Credentials) {
	// 初始化数据库连接
	dbConn = &db.DynamicSecretRotationDb{}
	err := dbConn.Init(&db.Config{
		DbConfig: &db.DbConfig{
			MaxOpenConns:        100,
			MaxIdleConns:        50,
			IdleTimeoutSeconds:  100,
			ReadTimeoutSeconds:  5,
			WriteTimeoutSeconds: 5,
			SecretName:          secretName, // 凭据名
			IpAddress:           dbAddress,  // 数据库地址
			Port:                dbPort,     // 数据库端口
			DbName:              dbName,     // 可以为空，或指定具体的数据库名
			ParamStr:            "charset=utf8&loc=Local",
		},
		SsmServiceConfig: &ssm.SsmAccount{
			SecretId:  credential.TmpSecretId,  // 需填写实际可用的SecretId
			SecretKey: credential.TmpSecretKey, // 需填写实际可用的SecretKey
			Token:     credential.Token,
			Region:    ssmRegionName, // 选择凭据所存储的地域
		},
		WatchChangeInterval: time.Second * 10, // 多长时间检查一下 凭据是否发生了轮转
	})
	if err != nil {
		fmt.Errorf("failed to init dbConn, err:%v\\n", err)
		return
	}
	// 模拟业务处理中，每过一段时间（一般是几毫秒），需要拿到db连接，来操作数据库的场景
	t := time.Tick(time.Second)
	for {
		select {
		case <-t:
			accessDb()
			queryDb()
		}
	}
}

func accessDb() {
	fmt.Println("--- accessDb start")
	c := dbConn.GetConn()
	if err := c.Ping(); err != nil {
		log.Fatal("failed to access db with err:", err)
	}
	log.Println("--- succeed to access db")
}

func queryDb() {
	var (
		id   int
		name string
	)
	log.Println("--- queryDb start")
	c := dbConn.GetConn()
	rows, err := c.Query("select id, name from user where id = ?", 1)
	if err != nil {
		log.Printf("failed to query db with err: ", err)
		log.Fatal(err)
	}
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&id, &name)
		if err != nil {
			log.Fatal(err)
		}
		log.Println(id, name)
	}
	err = rows.Err()
	if err != nil {
		log.Fatal(err)
	}
	log.Println("--- succeed to query db")
}