package main

import (
    "crypto/rand"
    "fmt"
    "math/big"
    "time"
)

func main() {
    for {
        fmt.Print("****************************************Rabin密码" +
            "****************************************\n 1.密钥生成 2.加密 3.解密 0.退出\n请输入操作编号：")
        var flag int
        fmt.Scanf("%d", &flag)
        if flag == 0 {
            return
        }
        if flag == 1 {
            fmt.Print("请输入要生成的密钥长度len（私钥长度为len,公钥长度为2len,单位bit）:")
            var len int
            fmt.Scanf("%d", &len)
            startTime := time.Now()
            p, q, n := genKey(len)
            fmt.Printf("公钥：%v\n私钥：%v和%v\n", n, p, q)
            fmt.Printf("耗时%f 毫秒\n", float64(time.Now().Sub(startTime).Microseconds())/1000.0)
        }
        if flag == 2 {
            var m, ns string
            n := new(big.Int)
            fmt.Print("请输入要加密的内容：")
            fmt.Scanf("%s", &m)
            fmt.Print("请输入公钥：")
            fmt.Scanf("%s", &ns)
            n.SetString(ns, 10)
            startTime := time.Now()
            fmt.Printf("密文：%v\n", enc(n, m))
            fmt.Printf("耗时%f 毫秒\n", float64(time.Now().Sub(startTime).Microseconds())/1000.0)

        }
        if flag == 3 {
            var ps, qs, c string
            p := new(big.Int)
            q := new(big.Int)
            fmt.Print("请输入要解密的内容：")
            fmt.Scanf("%s", &c)
            fmt.Print("请输入私钥p：")
            fmt.Scanf("%s", &ps)
            fmt.Print("请输入私钥q：")
            fmt.Scanf("%s", &qs)
            p.SetString(ps, 10)
            q.SetString(qs, 10)
            startTime := time.Now()
            list := dec(p, q, c)
            fmt.Printf("明文1：%v\n", list[0])
            fmt.Printf("明文2：%v\n", list[1])
            fmt.Printf("明文3：%v\n", list[2])
            fmt.Printf("明文4：%v\n", list[3])
            fmt.Printf("耗时%f 毫秒\n", float64(time.Now().Sub(startTime).Microseconds())/1000.0)

        }
    }
}
func enc(n *big.Int, m string) string {
    mb := new(big.Int)
    mb.SetBytes([]byte(m))
    c := new(big.Int)
    c = c.Mod(mb.Mul(mb, mb), n)
    return c.String()
}
func dec(p *big.Int, q *big.Int, c string) []string {
    n := new(big.Int).Mul(p, q)
    s := new(big.Int)
    t := new(big.Int)
    u := new(big.Int)
    v := new(big.Int)
    cb := new(big.Int)
    cb.SetString(c, 10)
    new(big.Int).GCD(s, t, p, q)
    u.Exp(cb, new(big.Int).Quo(new(big.Int).Add(p, big.NewInt(1)), big.NewInt(4)), p)
    v.Exp(cb, new(big.Int).Quo(new(big.Int).Add(q, big.NewInt(1)), big.NewInt(4)), q)
    m1 := new(big.Int).Mod(new(big.Int).Add(new(big.Int).Mul(new(big.Int).Mul(u, t), q), new(big.Int).Mul(new(big.Int).Mul(v, s), p)), n)
    m2 := new(big.Int).Mod(new(big.Int).Sub(new(big.Int).Mul(new(big.Int).Mul(u, t), q), new(big.Int).Mul(new(big.Int).Mul(v, s), p)), n)
    m3 := new(big.Int).Sub(n, m1)
    m4 := new(big.Int).Sub(n, m2)
    var mList []string
    mList = append(mList, string(m1.Bytes()))
    mList = append(mList, string(m2.Bytes()))
    mList = append(mList, string(m3.Bytes()))
    mList = append(mList, string(m4.Bytes()))
    return mList
}
func genKey(len int) (*big.Int, *big.Int, *big.Int) { //rabin密钥生成
    p := new(big.Int)
    q := new(big.Int)
    n := new(big.Int)
    for { //随机生成一个len比特的素数 并判断是否与3mod4同余
        t, _ := rand.Prime(rand.Reader, len)
        temp := new(big.Int)
        temp.Mod(t, big.NewInt(4))
        if big.NewInt(3).Cmp(temp) == 0 {
            p = t
            break
        }
    }
    for { //随机生成一个len比特的素数 并判断是否与3mod4同余 且不能和p相同
        t, _ := rand.Prime(rand.Reader, len)
        temp := new(big.Int)
        temp.Mod(t, big.NewInt(4))
        if big.NewInt(3).Cmp(temp) == 0 && t.Cmp(p) != 0 {
            q = t
            break
        }
    }
    n.Mul(p, q)
    return p, q, n
}