/*
   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
// TODO: remember to write newline at end
func (m *Moduli) Marshal() ([]byte, error) {
	var b bytes.Buffer
	b.Write([]byte(header))
	for _, i := range m.Params {
		line, err := i.marshalEntry()
		if err != nil {
			return b.Bytes(), err
		} else {
			b.Write(line)
		}
	}
	return b.Bytes(), nil
}

// marshalEntry is used to parse a specific DH entry into the moduli.
func (m *ModuliEntry) marshalEntry() ([]byte, 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),
		string(m.Size),
		string(m.Generator),
		mod,
	)
	return []byte(s), nil
}

// Unmarshal writes the Moduli format into m from the /etc/ssh/moduli format in data.
func Unmarshal(data []byte, m Moduli) error {
	var lines []string
	var entries []ModuliEntry
	lines = strings.Split(string(data), "\n")
	for _, line := range lines {
		e := ModuliEntry{}
		if reSkipLine.MatchString(line) {
			continue
		}
		l := strings.Fields(line)
		if err := unmarshalEntry(l, e); err != nil {
			return err
		}
		entries = append(entries, e)
	}
	m.Params = entries
	return nil
}

func unmarshalEntry(line []string, m ModuliEntry) error {
	if len(line) != 7 {
		return errors.New("field count mismatch")
	}
	if m.Time, err = time.Parse(timeFormat, line[0]); err != nil {
		return err
	}
	// Numeric types. Cast to uint8. There's probably a better way to do this but golang's pretty ugly with this stuff no matter what.
	// Type, Tests, Trials, Size, Generator
	conv := [5]uint8{}
	for idx := 1; idx <= 5; idx++ {
		v := line[idx]
		newv, err := strconv.Atoi(v)
		if err != nil {
			return err
		}
		conv[idx-1] = uint8(newv)
	}
	m.Type = conv[0]
	m.Tests = conv[1]
	m.Trials = conv[2]
	m.Size = conv[3]
	m.Generator = conv[4]
	// And the modulus convert to big.Int.
	modb, err := hex.DecodeString(line[6])
	if err != nil {
		return err
	}
	m.Modulus = big.Int{}
	m.Modulus.SetBytes(modb)
	return nil
}