package diffiehellman

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

// Random number generator
var rnd = rand.New(rand.NewSource(time.Now().UnixNano()))

// PrivateKey generates a number greater than 1 and less than p
func PrivateKey(p *big.Int) *big.Int {
	ret := &big.Int{}
	return ret.Add(big.NewInt(2), ret.Rand(rnd, ret.Sub(p, big.NewInt(2))))
}

// PublicKey generates the public key for private key p
func PublicKey(pKey, p *big.Int, g int64) *big.Int {
	ret := big.NewInt(g)
	return ret.Exp(ret, pKey, p)
}

// NewPair generates a public/private key pair from given primes p & g
func NewPair(p *big.Int, g int64) (*big.Int, *big.Int) {
	priv := PrivateKey(p)
	return priv, PublicKey(priv, p, g)
}

// SecretKey returns a secret key given private & public keys and the prime p
func SecretKey(priv, pub, p *big.Int) *big.Int {
	ret := &big.Int{}
	return ret.Exp(pub, priv, p)
}