/*
   SSHSecure - a program to harden OpenSSH from defaults
   Copyright (C) 2020  Brent Saner

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package moduli

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"net/http"
	"time"

	"github.com/Luzifer/go-dhparam"
	"golang.org/x/crypto/sha3"
)

// NewModuli returns a Moduli populated with Entry items.
func NewModuli(usePreGen ...bool) (m *Moduli, err error) {

	var doPreGen bool

	m = new(Moduli)

	if usePreGen != nil {
		doPreGen = usePreGen[0]
	} else {
		doPreGen = false
	}

	if doPreGen {
		if err = GetPreGen(m); err != nil {
			return
		}
		// This may take a while.
		if err = m.Harden(); err != nil {
			return
		}
	} else {
		if err = Generate(m); err != nil {
			return
		}
	}

	return
}

// GetPreGen gets the pregenerated moduli from upstream mirror.
func GetPreGen(m *Moduli) (err error) {

	var b []byte
	var goodCksum []byte
	var resp *http.Response

	// get the pregenerated moduli
	resp, err = http.Get(pregenURL)
	if err != nil {
		return
	}
	if resp.StatusCode != http.StatusOK {
		err = errors.New(fmt.Sprintf("returned status code %v: %v", resp.StatusCode, resp.Status))
		return
	}
	defer resp.Body.Close()

	b = make([]byte, resp.ContentLength)
	if _, err = resp.Body.Read(b); err != nil {
		return
	}

	// and compare the SHA3-512 (NIST) checksum.
	s := sha3.New512()
	if _, err = s.Write(b); err != nil {
		return
	}
	goodCksum, err = hex.DecodeString(pregenCksum)
	if err != nil {
		return
	}

	// We just compare the bytestrings.
	if bytes.Compare(s.Sum(nil), goodCksum) != 0 {
		err = errors.New("checksums do not match")
		return
	}

	if err := Unmarshal(b, m); err != nil {
		return
	}

	return
}

// Generate generates new moduli with Entry items. It's more secure than using GetPreGen (LogJam), but takes a *LOT* longer.
func Generate(m *Moduli) (err error) {

	var dh *dhparam.DH

	for _, bitLen := range genBits {
		for _, generator := range genGenerators {

			var e Entry

			e = Entry{
				Time:      time.Now(),
				Size:      bitLen,
				Generator: uint8(generator),
				/*
					Type:      0,
					Tests:     0,
					Trials:    0,
				*/
			}

			if dh, err = dhparam.Generate(int(bitLen), generator, nil); err != nil {
				continue // TODO: log/print
			}

			// Check() applies big.Int.ProbablyPrime() (Miller-Rabin - 0x04 in ssh moduli - and Baillie-PSW test), so it's probably fine.
			if errs, ok := dh.Check(); !ok {
				_ = errs // TODO: log/print
				continue
			} else {
				e.Time = time.Now()
				// e.Type =
				// e.Tests =
				// e.Trials =
				e.Modulus = *dh.P

				// TODO: https://stackoverflow.com/questions/18499352/golang-concurrency-how-to-append-to-the-same-slice-from-different-goroutines
				m.Groups = append(m.Groups, e)
			}

		}
	}

	return
}