package main

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// Snowflake 雪花算法生成器
type Snowflake struct {
	mu        sync.Mutex // 保证并发安全
	epoch     int64      // 起始时间戳（毫秒）
	nodeID    int64      // 节点ID
	sequence  int64      // 毫秒内序列号
	lastStamp int64      // 上一次生成ID的时间戳（毫秒）

	// 位段掩码与偏移量（根据64位结构定义）
	nodeIDBits  int64 = 10                     // 节点ID位数
	sequenceBits int64 = 12                    // 序列号位数
	nodeIDMax   int64 = -1 ^ (-1 << nodeIDBits) // 节点ID最大值（1023）
	sequenceMax int64 = -1 ^ (-1 << sequenceBits) // 序列号最大值（4095）

	nodeIDShift  int64 = sequenceBits          // 节点ID偏移量
	timeStampShift int64 = sequenceBits + nodeIDBits // 时间戳偏移量
}

// NewSnowflake 创建雪花算法实例
// epoch: 起始时间戳（毫秒），如 1704067200000（2024-01-01 00:00:00）
// nodeID: 节点ID（0-1023）
func NewSnowflake(epoch int64, nodeID int64) (*Snowflake, error) {
	// 校验节点ID范围
	if nodeID < 0 || nodeID > (-1 ^ (-1 << 10)) {
		return nil, errors.New("node ID must be between 0 and 1023")
	}

	return &Snowflake{
		epoch:     epoch,
		nodeID:    nodeID,
		sequence:  0,
		lastStamp: -1,
	}, nil
}

// 获取当前毫秒级时间戳
func (s *Snowflake) now() int64 {
	return time.Now().UnixMilli()
}

// NextID 生成下一个唯一ID
func (s *Snowflake) NextID() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 1. 获取当前时间戳
	now := s.now()

	// 2. 处理时钟回拨
	if now < s.lastStamp {
		return 0, fmt.Errorf("clock moved backwards. Refusing to generate ID for %d milliseconds", s.lastStamp-now)
	}

	// 3. 同一毫秒内，序列号自增
	if now == s.lastStamp {
		s.sequence = (s.sequence + 1) & s.sequenceMax
		// 序列号溢出（同一毫秒生成超过4096个ID）
		if s.sequence == 0 {
			// 等待到下一毫秒
			for now <= s.lastStamp {
				now = s.now()
			}
		}
	} else {
		// 不同毫秒，重置序列号
		s.sequence = 0
	}

	// 4. 更新上一次时间戳
	s.lastStamp = now

	// 5. 组合64位ID
	// 时间戳部分：(now - epoch) << 22
	// 节点ID部分：nodeID << 12
	// 序列号部分：sequence
	id := ((now - s.epoch) << s.timeStampShift) | (s.nodeID << s.nodeIDShift) | s.sequence

	return id, nil
}

// 测试示例
func main() {
	// 自定义起始时间（2024-01-01 00:00:00）
	epoch := int64(1704067200000)
	// 节点ID（假设当前节点为1）
	nodeID := int64(1)

	// 创建雪花算法实例
	sf, err := NewSnowflake(epoch, nodeID)
	if err != nil {
		panic(err)
	}

	// 并发生成10个ID测试
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			id, err := sf.NextID()
			if err != nil {
				fmt.Printf("生成ID失败：%v\n", err)
				return
			}
			fmt.Printf("第%d个ID：%d\n", idx+1, id)
		}(i)
	}
	wg.Wait()
}

package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"fmt"
)

// PKCS7 补码（AES 要求明文长度为 16 字节倍数）
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS7 解补码
func pkcs7UnPadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, fmt.Errorf("加密数据为空")
	}
	// 取出最后一个字节的值，即为补码长度
	padding := int(data[length-1])
	return data[:length-padding], nil
}

// AESCBCEncrypt AES-CBC 加密
// key: 密钥（16/24/32 字节，对应 AES-128/AES-192/AES-256）
// iv: 初始化向量（必须 16 字节）
func AESCBCEncrypt(plainText, key, iv []byte) (string, error) {
	// 1. 创建 AES 密码块
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	// 2. 补码
	plainText = pkcs7Padding(plainText, block.BlockSize())
	// 3. 创建 CBC 模式
	mode := cipher.NewCBCEncrypter(block, iv)
	// 4. 加密
	cipherText := make([]byte, len(plainText))
	mode.CryptBlocks(cipherText, plainText)
	// 5. base64 编码返回
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

// AESCBCDecrypt AES-CBC 解密
func AESCBCDecrypt(cipherTextBase64 string, key, iv []byte) (string, error) {
	// 1. base64 解码
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return "", err
	}
	// 2. 创建 AES 密码块
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	// 3. 创建 CBC 解密模式
	mode := cipher.NewCBCDecrypter(block, iv)
	// 4. 解密
	plainText := make([]byte, len(cipherText))
	mode.CryptBlocks(plainText, cipherText)
	// 5. 解补码
	plainText, err = pkcs7UnPadding(plainText)
	if err != nil {
		return "", err
	}
	return string(plainText), nil
}

func main() {
	// 测试 AES-CBC
	plainText := []byte("Hello, AES Encryption!")
	// 密钥（16 字节 = AES-128）
	key := []byte("1234567890123456")
	// IV（必须 16 字节）
	iv := []byte("1234567890123456")

	// 加密
	cipherText, err := AESCBCEncrypt(plainText, key, iv)
	if err != nil {
		fmt.Printf("AES 加密失败：%v\n", err)
		return
	}
	fmt.Printf("AES 加密结果：%s\n", cipherText)

	// 解密
	decryptedText, err := AESCBCDecrypt(cipherText, key, iv)
	if err != nil {
		fmt.Printf("AES 解密失败：%v\n", err)
		return
	}
	fmt.Printf("AES 解密结果：%s\n", decryptedText)
}

package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"os"
)

// GenerateRSAKey 生成 RSA 密钥对（2048 位）
func GenerateRSAKey() (pubKey []byte, priKey []byte, err error) {
	// 1. 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}
	// 2. 编码私钥（PEM 格式）
	privateBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateBytes,
	}
	priKey = pem.EncodeToMemory(privateBlock)

	// 3. 提取公钥
	publicKey := &privateKey.PublicKey
	// 4. 编码公钥（PEM 格式）
	publicBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return nil, nil, err
	}
	publicBlock := &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: publicBytes,
	}
	pubKey = pem.EncodeToMemory(publicBlock)

	return pubKey, priKey, nil
}

// RSAEncrypt 公钥加密
func RSAEncrypt(plainText []byte, pubKey []byte) (string, error) {
	// 1. 解析公钥
	block, _ := pem.Decode(pubKey)
	if block == nil {
		return "", fmt.Errorf("公钥解析失败")
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return "", err
	}
	publicKey, ok := pubInterface.(*rsa.PublicKey)
	if !ok {
		return "", fmt.Errorf("公钥类型错误")
	}

	// 2. 加密（OAEP 填充更安全，PKCS1v15 兼容老系统）
	cipherText, err := rsa.EncryptOAEP(
		rand.Reader,
		publicKey,
		nil,
		plainText,
		nil,
	)
	if err != nil {
		return "", err
	}

	// 3. base64 编码
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

// RSADecrypt 私钥解密
func RSADecrypt(cipherTextBase64 string, priKey []byte) (string, error) {
	// 1. base64 解码
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return "", err
	}

	// 2. 解析私钥
	block, _ := pem.Decode(priKey)
	if block == nil {
		return "", fmt.Errorf("私钥解析失败")
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	// 3. 解密
	plainText, err := rsa.DecryptOAEP(
		rand.Reader,
		privateKey,
		nil,
		cipherText,
		nil,
	)
	if err != nil {
		return "", err
	}

	return string(plainText), nil
}

func main() {
	// 生成 RSA 密钥对
	pubKey, priKey, err := GenerateRSAKey()
	if err != nil {
		fmt.Printf("生成 RSA 密钥失败：%v\n", err)
		return
	}
	fmt.Printf("RSA 公钥：\n%s\n", string(pubKey))
	fmt.Printf("RSA 私钥：\n%s\n", string(priKey))

	// 测试加密解密
	plainText := []byte("Hello, RSA Encryption!")
	// 公钥加密
	cipherText, err := RSAEncrypt(plainText, pubKey)
	if err != nil {
		fmt.Printf("RSA 加密失败：%v\n", err)
		return
	}
	fmt.Printf("RSA 加密结果：%s\n", cipherText)

	// 私钥解密
	decryptedText, err := RSADecrypt(cipherText, priKey)
	if err != nil {
		fmt.Printf("RSA 解密失败：%v\n", err)
		return
	}
	fmt.Printf("RSA 解密结果：%s\n", decryptedText)
}

package main

import (
	"encoding/base64"
	"fmt"

	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm4"
	"github.com/tjfoc/gmsm/x509"
)

// ===================== SM4 对称加密 =====================
// SM4ECBEncrypt SM4-ECB 加密（ECB 模式，无需 IV）
func SM4ECBEncrypt(plainText, key []byte) (string, error) {
	// 1. 创建 SM4 密码块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}
	// 2. 补码（PKCS7，与 AES 通用）
	plainText = pkcs7Padding(plainText, block.BlockSize())
	// 3. 加密（ECB 模式）
	cipherText := make([]byte, len(plainText))
	for i := 0; i < len(plainText); i += block.BlockSize() {
		block.Encrypt(cipherText[i:i+block.BlockSize()], plainText[i:i+block.BlockSize()])
	}
	// 4. base64 编码
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

// SM4ECBDecrypt SM4-ECB 解密
func SM4ECBDecrypt(cipherTextBase64 string, key []byte) (string, error) {
	// 1. base64 解码
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return "", err
	}
	// 2. 创建 SM4 密码块
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}
	// 3. 解密
	plainText := make([]byte, len(cipherText))
	for i := 0; i < len(cipherText); i += block.BlockSize() {
		block.Decrypt(plainText[i:i+block.BlockSize()], cipherText[i:i+block.BlockSize()])
	}
	// 4. 解补码
	plainText, err = pkcs7UnPadding(plainText)
	if err != nil {
		return "", err
	}
	return string(plainText), nil
}

// ===================== SM2 非对称加密 =====================
// GenerateSM2Key 生成 SM2 密钥对
func GenerateSM2Key() (pubKey, priKey []byte, err error) {
	// 1. 生成 SM2 私钥
	privateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	// 2. 编码私钥（PEM 格式）
	privateBytes, err := x509.WritePrivateKeyToPem(privateKey, nil)
	if err != nil {
		return nil, nil, err
	}
	// 3. 编码公钥（PEM 格式）
	publicBytes, err := x509.WritePublicKeyToPem(&privateKey.PublicKey)
	if err != nil {
		return nil, nil, err
	}
	return publicBytes, privateBytes, nil
}

// SM2Encrypt SM2 公钥加密
func SM2Encrypt(plainText []byte, pubKey []byte) (string, error) {
	// 1. 解析公钥
	publicKey, err := x509.ReadPublicKeyFromPem(pubKey)
	if err != nil {
		return "", err
	}
	// 2. 加密（SM2 推荐 C1C3C2 格式）
	cipherText, err := sm2.Encrypt(publicKey, plainText, rand.Reader, sm2.C1C3C2)
	if err != nil {
		return "", err
	}
	// 3. base64 编码
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

// SM2Decrypt SM2 私钥解密
func SM2Decrypt(cipherTextBase64 string, priKey []byte) (string, error) {
	// 1. base64 解码
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return "", err
	}
	// 2. 解析私钥
	privateKey, err := x509.ReadPrivateKeyFromPem(priKey, nil)
	if err != nil {
		return "", err
	}
	// 3. 解密
	plainText, err := sm2.Decrypt(privateKey, cipherText, sm2.C1C3C2)
	if err != nil {
		return "", err
	}
	return string(plainText), nil
}

func main() {
	// ===================== 测试 SM4 =====================
	sm4Key := []byte("1234567890123456") // SM4 密钥必须 16 字节
	sm4PlainText := []byte("Hello, SM4 Encryption!")
	// SM4 加密
	sm4CipherText, err := SM4ECBEncrypt(sm4PlainText, sm4Key)
	if err != nil {
		fmt.Printf("SM4 加密失败：%v\n", err)
		return
	}
	fmt.Printf("SM4 加密结果：%s\n", sm4CipherText)
	// SM4 解密
	sm4DecryptedText, err := SM4ECBDecrypt(sm4CipherText, sm4Key)
	if err != nil {
		fmt.Printf("SM4 解密失败：%v\n", err)
		return
	}
	fmt.Printf("SM4 解密结果：%s\n", sm4DecryptedText)

	// ===================== 测试 SM2 =====================
	// 生成 SM2 密钥对
	sm2PubKey, sm2PriKey, err := GenerateSM2Key()
	if err != nil {
		fmt.Printf("生成 SM2 密钥失败：%v\n", err)
		return
	}
	sm2PlainText := []byte("Hello, SM2 Encryption!")
	// SM2 加密
	sm2CipherText, err := SM2Encrypt(sm2PlainText, sm2PubKey)
	if err != nil {
		fmt.Printf("SM2 加密失败：%v\n", err)
		return
	}
	fmt.Printf("SM2 加密结果：%s\n", sm2CipherText)
	// SM2 解密
	sm2DecryptedText, err := SM2Decrypt(sm2CipherText, sm2PriKey)
	if err != nil {
		fmt.Printf("SM2 解密失败：%v\n", err)
		return
	}
	fmt.Printf("SM2 解密结果：%s\n", sm2DecryptedText)
}