path: root/lib/rsa.go

// Copyright © 2021 siddharth ravikumar <s@ricketyspace.net>
// SPDX-License-Identifier: ISC

package lib

import (
	"crypto/rand"
	"errors"
	"math/big"
)

// Represents an RSA key pair.
type RSAPair struct {
	Public  *RSAPub
	Private *RSAPrivate
}

type RSAPub struct {
	e *big.Int
	n *big.Int
}

type RSAPrivate struct {
	d *big.Int
	n *big.Int
}

// Copy b to a.
func biCopy(a, b *big.Int) *big.Int {
	a.SetBytes(b.Bytes())
	if b.Sign() == -1 {
		a.Mul(a, big.NewInt(-1))
	}
	return a
}

func InvMod(a, n *big.Int) (*big.Int, error) {
	// Initialize.
	t0 := big.NewInt(0)
	t1 := big.NewInt(1)
	r0 := biCopy(big.NewInt(0), n)
	r1 := biCopy(big.NewInt(0), a)

	for r1.Cmp(big.NewInt(0)) != 0 {
		q := big.NewInt(0)
		q.Div(r0, r1)

		tt := big.NewInt(0)
		tt = tt.Mul(q, t1)
		tt = tt.Sub(t0, tt)

		biCopy(t0, t1)
		biCopy(t1, tt)

		tr := big.NewInt(0)
		tr = tr.Mul(q, r1)
		tr = tr.Sub(r0, tr)

		biCopy(r0, r1)
		biCopy(r1, tr)
	}

	if r0.Cmp(big.NewInt(1)) > 0 {
		return nil, CPError{"not invertible"}
	}
	if t0.Cmp(big.NewInt(0)) < 0 {
		t0.Add(t0, n)
	}
	return t0, nil
}

// RSAGenKey generates a rsa key pair with N equal to size bits.
func RSAGenKey(size int) (*RSAPair, error) {
	// Initialize.
	e := big.NewInt(3)
	d := big.NewInt(0)
	n := big.NewInt(0)

	// Compute n and d.
	for {
		// Generate prime p.
		p, err := rand.Prime(rand.Reader, size/2)
		if err != nil {
			return nil, CPError{"unable to generate p"}
		}

		// Generate prime q.
		q, err := rand.Prime(rand.Reader, size/2)
		if err != nil {
			return nil, CPError{"unable to generate q"}
		}

		// Calculate n.
		n = big.NewInt(0).Mul(p, q)

		// Calculate totient.
		p1 := big.NewInt(0).Sub(p, big.NewInt(1)) // p-1
		q1 := big.NewInt(0).Sub(q, big.NewInt(1)) // q-1
		et := big.NewInt(0).Mul(p1, q1)           // Totient `et`.

		// Calculate private key `d`.
		d, err = InvMod(e, et)
		if err != nil {
			continue // Inverse does not does. Try again.
		}
		break
	}
	if n.Cmp(big.NewInt(0)) <= 0 {
		return nil, CPError{"unable to compute n"}
	}
	if d.Cmp(big.NewInt(0)) <= 0 {
		return nil, CPError{"unable to compute d"}
	}

	// Make pub key.
	pub := new(RSAPub)
	pub.e = e
	pub.n = biCopy(big.NewInt(0), n)

	// Make private key.
	prv := new(RSAPrivate)
	prv.d = d
	prv.n = biCopy(big.NewInt(0), n)

	// Make key pair.
	pair := new(RSAPair)
	pair.Public = pub
	pair.Private = prv

	return pair, nil
}

func (r *RSAPub) Encrypt(msg []byte) []byte {
	// Convert message to big int.
	m := big.NewInt(0).SetBytes(msg)

	// Encrypt.
	c := big.NewInt(0).Exp(m, r.e, r.n)

	return c.Bytes()
}

// LazyVerify verifies the RSA signature sig with the given message
// msg.
//
// Returns true if the verification is successful; false otherwise.
func (r *RSAPub) LazyVerify(msg, sig []byte) bool {
	// Hash message using md4.
	var (
		dgst Md4    = Md4{}
		vd   []byte // Data to verify signature with.
	)
	dgst.Init([]uint32{})
	dgst.Message(msg)
	vd = dgst.Hash()

	var (
		y  *big.Int // Signature as in integer.
		x  *big.Int // Encryption block as an integer.
		eb []byte   // Encryption block.
	)

	// Convert signature to integer.
	y = new(big.Int).SetBytes(sig)

	// Get encryption block.
	x = new(big.Int).Exp(y, r.e, r.n)
	eb = x.Bytes()

	if eb[0] != 0x01 {
		return false
	}
	if eb[1] != 0xFF {
		return false
	}
	var (
		ffd bool
		lb  byte = eb[0]
		d   []byte
	)
	for _, b := range eb[1:] {
		if ffd {
			d = append(d, b)
		}
		if lb == 0xFF && b == 0x00 {
			ffd = true
		}
		lb = b
	}
	return BytesEqual(vd, d[:len(vd)])
}

func (r *RSAPub) E() *big.Int {
	return r.e
}

func (r *RSAPub) N() *big.Int {
	return r.n
}

func (r *RSAPrivate) Decrypt(cipher []byte) []byte {
	// Convert cipher to big int.
	c := big.NewInt(0).SetBytes(cipher)

	// Decrypt.
	m := big.NewInt(0).Exp(c, r.d, r.n)

	return m.Bytes()
}

// Sign generates a signature for the give data d according to RFC2313
// sections 8 and 10.
//
// https://datatracker.ietf.org/doc/html/rfc2313
func (r *RSAPrivate) Sign(msg []byte) (ed []byte, err error) {
	// Hash message using md4.
	var (
		dgst Md4    = Md4{}
		d    []byte // Data to sign.
	)
	dgst.Init([]uint32{})
	dgst.Message(msg)
	d = dgst.Hash()

	var (
		kn int = r.n.BitLen() / 8 // Size of N in bytes.
		dn int = len(d)           // Length of d.
		pn int = kn - 3 - dn      // Number of padding characters.

		eb []byte   // Encryption block.
		x  *big.Int // Encryption block as an integer.
		y  *big.Int // Signature as an integer.
	)
	if dn > kn-11 {
		err = errors.New("rsa: data len > k-11")
		return
	}

	// Build encryption block.
	eb = []byte{0x00, 0x01}
	for i := 0; i < pn; i++ {
		eb = append(eb, 0xFF)
	}
	eb = append(eb, 0x00)
	eb = append(eb, d...)

	// Convert to big int.
	x = new(big.Int).SetBytes(eb)

	// Sign.
	y = big.NewInt(0).Exp(x, r.d, r.n)

	ed = y.Bytes()
	return
}