/*
   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"
	"math/big"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var reSkipLine, _ = regexp.Compile(`^\s*(#.*)?$`)

// Marshal returns the /etc/ssh/moduli format of m.
// Format of: Time Type Tests Tries Size Generator Modulus
func (m *Moduli) Marshal() (bytesOut []byte, err error) {

	var b bytes.Buffer

	b.Write([]byte(header))
	for _, i := range m.Groups {
		line, err := i.marshalEntry()
		if err != nil {
			return nil, err
		} else {
			b.Write(line)
		}
	}

	b.Write([]byte("\n"))

	bytesOut = b.Bytes()

	return
}

// marshalEntry is used to parse an Entry into the moduli(5) format.
func (m *Entry) marshalEntry() (sum []byte, err error) {

	mod := hex.EncodeToString(m.Modulus.Bytes())

	s := fmt.Sprintf(
		"%v %v %v %v %v %v %v\n",
		m.Time.Format(timeFormat),
		string(m.Type),
		string(m.Tests),
		string(m.Trials),
		strconv.Itoa(int(m.Size)-1), // see this thread https://twitter.com/SysAdm_Podcast/status/1386714803679399940
		string(m.Generator),
		mod,
	)

	sum = []byte(s)

	return
}

// Unmarshal populates a Moduli from the /etc/ssh/moduli format.
func Unmarshal(data []byte, m *Moduli) (err error) {

	var lines []string
	var entries []Entry

	lines = strings.Split(string(data), "\n")
	for _, line := range lines {

		e := Entry{}

		if reSkipLine.MatchString(line) {
			continue
		}

		l := strings.Fields(line)
		if err := unmarshalEntry(l, e); err != nil {
			return
		}

		entries = append(entries, e)
	}

	m.Groups = entries

	return
}

// unmarshalEntry unmarshals a single line from an /etc/ssh/moduli into an Entry.
func unmarshalEntry(line []string, m Entry) (err error) {

	var modb []byte

	if len(line) != 7 {
		err = errors.New("field count mismatch")
		return
	}

	if m.Time, err = time.Parse(timeFormat, line[0]); err != nil {
		return
	}
	// Numeric types. Cast to uint16. There's probably a better way to do this but golang's pretty ugly with this stuff no matter what.
	// The worst part is all of them are uint8 except size (uint16).
	// Type, Tests, Trials, Size, Generator
	conv := [5]uint16{}
	for idx := 1; idx <= 5; idx++ {
		v := line[idx]
		newv, err := strconv.Atoi(v)
		if err != nil {
			return err
		}
		conv[idx-1] = uint16(newv)
	}

	m.Type = uint8(conv[0])
	m.Tests = uint8(conv[1])
	m.Trials = uint8(conv[2])
	m.Size = conv[3] + 1 // see this thread https://twitter.com/SysAdm_Podcast/status/1386714803679399940
	m.Generator = uint8(conv[4])

	// And the modulus convert to big.Int.
	if modb, err = hex.DecodeString(line[6]); err != nil {
		return
	}
	m.Modulus = big.Int{}
	m.Modulus.SetBytes(modb)

	return
}

func (m *Moduli) Harden() (err error) {

	var entries []Entry

	for _, e := range m.Groups {

		e.Time = time.Now()

		if e.Size >= minBits {
			entries = append(entries, e)
		}
	}
	m.Groups = entries

	if len(m.Groups) < recMinMod {
		err = errors.New("does not meet recommended minimum moduli")
		return
	}

	// TODO: find way of testing/sieving primes

	return
}