// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package dsa implements the Digital Signature Algorithm, as defined in FIPS 186-3
package dsa

import (
	"big"
	"errors"
	"io"
)

// Parameters represents the domain parameters for a key. These parameters can
// be shared across many keys. The bit length of Q must be a multiple of 8.
type Parameters struct {
	P, Q, G *big.Int
}

// PublicKey represents a DSA public key.
type PublicKey struct {
	Parameters
	Y *big.Int
}

// PrivateKey represents a DSA private key.
type PrivateKey struct {
	PublicKey
	X *big.Int
}

type invalidPublicKeyError int

func (invalidPublicKeyError) Error() string {
	return "crypto/dsa: invalid public key"
}

// InvalidPublicKeyError results when a public key is not usable by this code.
// FIPS is quite strict about the format of DSA keys, but other code may be
// less so. Thus, when using keys which may have been generated by other code,
// this error must be handled.
var InvalidPublicKeyError = invalidPublicKeyError(0)

// ParameterSizes is a enumeration of the acceptable bit lengths of the primes
// in a set of DSA parameters. See FIPS 186-3, section 4.2.
type ParameterSizes int

const (
	L1024N160 ParameterSizes = iota
	L2048N224
	L2048N256
	L3072N256
)

// numMRTests is the number of Miller-Rabin primality tests that we perform. We
// pick the largest recommended number from table C.1 of FIPS 186-3.
const numMRTests = 64

// GenerateParameters puts a random, valid set of DSA parameters into params.
// This function takes many seconds, even on fast machines.
func GenerateParameters(params *Parameters, rand io.Reader, sizes ParameterSizes) (err error) {
	// This function doesn't follow FIPS 186-3 exactly in that it doesn't
	// use a verification seed to generate the primes. The verification
	// seed doesn't appear to be exported or used by other code and
	// omitting it makes the code cleaner.

	var L, N int
	switch sizes {
	case L1024N160:
		L = 1024
		N = 160
	case L2048N224:
		L = 2048
		N = 224
	case L2048N256:
		L = 2048
		N = 256
	case L3072N256:
		L = 3072
		N = 256
	default:
		return errors.New("crypto/dsa: invalid ParameterSizes")
	}

	qBytes := make([]byte, N/8)
	pBytes := make([]byte, L/8)

	q := new(big.Int)
	p := new(big.Int)
	rem := new(big.Int)
	one := new(big.Int)
	one.SetInt64(1)

GeneratePrimes:
	for {
		_, err = io.ReadFull(rand, qBytes)
		if err != nil {
			return
		}

		qBytes[len(qBytes)-1] |= 1
		qBytes[0] |= 0x80
		q.SetBytes(qBytes)

		if !big.ProbablyPrime(q, numMRTests) {
			continue
		}

		for i := 0; i < 4*L; i++ {
			_, err = io.ReadFull(rand, pBytes)
			if err != nil {
				return
			}

			pBytes[len(pBytes)-1] |= 1
			pBytes[0] |= 0x80

			p.SetBytes(pBytes)
			rem.Mod(p, q)
			rem.Sub(rem, one)
			p.Sub(p, rem)
			if p.BitLen() < L {
				continue
			}

			if !big.ProbablyPrime(p, numMRTests) {
				continue
			}

			params.P = p
			params.Q = q
			break GeneratePrimes
		}
	}

	h := new(big.Int)
	h.SetInt64(2)
	g := new(big.Int)

	pm1 := new(big.Int).Sub(p, one)
	e := new(big.Int).Div(pm1, q)

	for {
		g.Exp(h, e, p)
		if g.Cmp(one) == 0 {
			h.Add(h, one)
			continue
		}

		params.G = g
		return
	}

	panic("unreachable")
}

// GenerateKey generates a public&private key pair. The Parameters of the
// PrivateKey must already be valid (see GenerateParameters).
func GenerateKey(priv *PrivateKey, rand io.Reader) error {
	if priv.P == nil || priv.Q == nil || priv.G == nil {
		return errors.New("crypto/dsa: parameters not set up before generating key")
	}

	x := new(big.Int)
	xBytes := make([]byte, priv.Q.BitLen()/8)

	for {
		_, err := io.ReadFull(rand, xBytes)
		if err != nil {
			return err
		}
		x.SetBytes(xBytes)
		if x.Sign() != 0 && x.Cmp(priv.Q) < 0 {
			break
		}
	}

	priv.X = x
	priv.Y = new(big.Int)
	priv.Y.Exp(priv.G, x, priv.P)
	return nil
}

// Sign signs an arbitrary length hash (which should be the result of hashing a
// larger message) using the private key, priv. It returns the signature as a
// pair of integers. The security of the private key depends on the entropy of
// rand.
func Sign(rand io.Reader, priv *PrivateKey, hash []byte) (r, s *big.Int, err error) {
	// FIPS 186-3, section 4.6

	n := priv.Q.BitLen()
	if n&7 != 0 {
		err = InvalidPublicKeyError
		return
	}
	n >>= 3

	for {
		k := new(big.Int)
		buf := make([]byte, n)
		for {
			_, err = io.ReadFull(rand, buf)
			if err != nil {
				return
			}
			k.SetBytes(buf)
			if k.Sign() > 0 && k.Cmp(priv.Q) < 0 {
				break
			}
		}

		kInv := new(big.Int).ModInverse(k, priv.Q)

		r = new(big.Int).Exp(priv.G, k, priv.P)
		r.Mod(r, priv.Q)

		if r.Sign() == 0 {
			continue
		}

		if n > len(hash) {
			n = len(hash)
		}
		z := k.SetBytes(hash[:n])

		s = new(big.Int).Mul(priv.X, r)
		s.Add(s, z)
		s.Mod(s, priv.Q)
		s.Mul(s, kInv)
		s.Mod(s, priv.Q)

		if s.Sign() != 0 {
			break
		}
	}

	return
}

// Verify verifies the signature in r, s of hash using the public key, pub. It
// returns true iff the signature is valid.
func Verify(pub *PublicKey, hash []byte, r, s *big.Int) bool {
	// FIPS 186-3, section 4.7

	if r.Sign() < 1 || r.Cmp(pub.Q) >= 0 {
		return false
	}
	if s.Sign() < 1 || s.Cmp(pub.Q) >= 0 {
		return false
	}

	w := new(big.Int).ModInverse(s, pub.Q)

	n := pub.Q.BitLen()
	if n&7 != 0 {
		return false
	}
	n >>= 3

	if n > len(hash) {
		n = len(hash)
	}
	z := new(big.Int).SetBytes(hash[:n])

	u1 := new(big.Int).Mul(z, w)
	u1.Mod(u1, pub.Q)
	u2 := w.Mul(r, w)
	u2.Mod(u2, pub.Q)
	v := u1.Exp(pub.G, u1, pub.P)
	u2.Exp(pub.Y, u2, pub.P)
	v.Mul(v, u2)
	v.Mod(v, pub.P)
	v.Mod(v, pub.Q)

	return v.Cmp(r) == 0
}