// This file is part of the happyDomain (R) project.
// Copyright (c) 2020-2026 happyDomain
// Authors: Pierre-Olivier Mercier, et al.
//
// This program is offered under a commercial and under the AGPL license.
// For commercial licensing, contact us at <contact@happydomain.org>.
//
// For AGPL licensing:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package checker

import (
	"bufio"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"strings"
)

// The SSH transport protocol is fully specified in RFC 4253. For the
// security audit we only need the pre-authentication handshake:
//
//   1. exchange of protocol version strings (ASCII banners, CRLF-terminated)
//   2. exchange of SSH_MSG_KEXINIT packets, which carry the full algorithm
//      preference lists in one go.
//
// We deliberately avoid going beyond KEXINIT: a shallow probe is cheap,
// works against every RFC-compliant SSH server without depending on any
// particular algorithm family being supported, and sidesteps the risk of
// running actual KEX math with untrusted peers.

const (
	// sshClientBanner is the version string we advertise. ssh-audit and
	// nmap publish a recognisable banner so operators can tell an audit
	// probe apart from a real client in their logs.
	sshClientBanner = "SSH-2.0-happyDomain-checker_1.0"

	msgKexInit = 20

	// maxPacketSize caps the largest packet we will read. RFC 4253 allows
	// up to 32768 bytes of payload, so 65k is a safe ceiling that also
	// protects us from a rogue server trying to exhaust memory.
	maxPacketSize = 65535

	// maxBannerSize limits how much we read before we give up on the
	// peer's version string. RFC 4253 mandates at most 255 bytes.
	maxBannerSize = 4096
)

// kexInitPayload is the parsed contents of a KEXINIT packet. The field
// names match RFC 4253 §7.1 verbatim to make audits easy.
type kexInitPayload struct {
	Cookie                          [16]byte
	KexAlgorithms                   []string
	ServerHostKeyAlgorithms         []string
	EncryptionAlgorithmsClientToSvr []string
	EncryptionAlgorithmsSvrToClient []string
	MACAlgorithmsClientToSvr        []string
	MACAlgorithmsSvrToClient        []string
	CompressionAlgorithmsClientToSv []string
	CompressionAlgorithmsSvrToClt   []string
	LanguagesClientToSvr            []string
	LanguagesSvrToClient            []string
	FirstKexPacketFollows           bool
}

// readBanner reads the peer's SSH identification string. Servers may
// send several CRLF-terminated lines of free-text before the actual
// "SSH-2.0-..." line (RFC 4253 §4.2); we skip those and return the first
// line that looks like a version exchange.
func readBanner(r *bufio.Reader) (string, error) {
	for i := 0; i < 16; i++ {
		line, err := readLine(r, maxBannerSize)
		if err != nil {
			return "", err
		}
		if strings.HasPrefix(line, "SSH-") {
			return line, nil
		}
	}
	return "", fmt.Errorf("no SSH version string received")
}

// readLine reads a single CRLF-terminated line (or LF-terminated, as
// some servers omit the CR) and returns it without the terminator.
func readLine(r *bufio.Reader, max int) (string, error) {
	var buf []byte
	for {
		b, err := r.ReadByte()
		if err != nil {
			return "", err
		}
		if b == '\n' {
			if n := len(buf); n > 0 && buf[n-1] == '\r' {
				buf = buf[:n-1]
			}
			return string(buf), nil
		}
		buf = append(buf, b)
		if len(buf) > max {
			return "", fmt.Errorf("line too long")
		}
	}
}

// writeBanner sends our client identification string.
func writeBanner(w io.Writer) error {
	_, err := io.WriteString(w, sshClientBanner+"\r\n")
	return err
}

// readPacket reads a single SSH binary packet (RFC 4253 §6) from r. The
// handshake is still in cleartext at this point, so we don't worry
// about MAC or cipher state: packet_length + padding_length + payload +
// random padding, no MAC.
func readPacket(r io.Reader) (payload []byte, err error) {
	var lenBuf [4]byte
	if _, err = io.ReadFull(r, lenBuf[:]); err != nil {
		return nil, err
	}
	packetLen := binary.BigEndian.Uint32(lenBuf[:])
	if packetLen < 5 || packetLen > maxPacketSize {
		return nil, fmt.Errorf("invalid packet length %d", packetLen)
	}
	body := make([]byte, packetLen)
	if _, err = io.ReadFull(r, body); err != nil {
		return nil, err
	}
	padLen := int(body[0])
	if padLen >= len(body) {
		return nil, fmt.Errorf("invalid padding length %d vs packet %d", padLen, len(body))
	}
	return body[1 : len(body)-padLen], nil
}

// writePacket frames payload into an RFC 4253 binary packet and sends it.
func writePacket(w io.Writer, payload []byte) error {
	// packet_length covers padding_length + payload + random_padding,
	// but not itself. The total (4 + packet_length) must be a multiple
	// of 8 (the block size used in unencrypted mode), and padding must
	// be at least 4 bytes.
	const blockSize = 8
	padLen := blockSize - ((5 + len(payload)) % blockSize)
	if padLen < 4 {
		padLen += blockSize
	}
	packetLen := 1 + len(payload) + padLen

	buf := make([]byte, 4+packetLen)
	binary.BigEndian.PutUint32(buf[:4], uint32(packetLen))
	buf[4] = byte(padLen)
	copy(buf[5:], payload)
	if _, err := rand.Read(buf[5+len(payload):]); err != nil {
		return fmt.Errorf("padding rand: %w", err)
	}
	_, err := w.Write(buf)
	return err
}

// buildKexInit crafts a client KEXINIT payload that advertises every
// algorithm family the Go SSH stack knows, plus the typical OpenSSH
// names we aren't implementing. We are never going to actually perform
// key exchange over this connection: the server only needs to accept
// our KEXINIT as well-formed and echo its own.
func buildKexInit() []byte {
	var cookie [16]byte
	_ = mustRand(cookie[:])

	kex := strings.Join([]string{
		"curve25519-sha256",
		"curve25519-sha256@libssh.org",
		"ecdh-sha2-nistp256",
		"ecdh-sha2-nistp384",
		"ecdh-sha2-nistp521",
		"diffie-hellman-group-exchange-sha256",
		"diffie-hellman-group16-sha512",
		"diffie-hellman-group14-sha256",
		"diffie-hellman-group14-sha1",
		"diffie-hellman-group1-sha1",
		"sntrup761x25519-sha512@openssh.com",
		"mlkem768x25519-sha256",
	}, ",")
	hostKey := strings.Join([]string{
		"ssh-ed25519",
		"ssh-ed25519-cert-v01@openssh.com",
		"rsa-sha2-512",
		"rsa-sha2-256",
		"ssh-rsa",
		"ecdsa-sha2-nistp256",
		"ecdsa-sha2-nistp384",
		"ecdsa-sha2-nistp521",
		"ssh-dss",
	}, ",")
	ciphers := strings.Join([]string{
		"chacha20-poly1305@openssh.com",
		"aes256-gcm@openssh.com",
		"aes128-gcm@openssh.com",
		"aes256-ctr",
		"aes192-ctr",
		"aes128-ctr",
		"aes256-cbc",
		"aes128-cbc",
		"3des-cbc",
	}, ",")
	macs := strings.Join([]string{
		"hmac-sha2-512-etm@openssh.com",
		"hmac-sha2-256-etm@openssh.com",
		"umac-128-etm@openssh.com",
		"hmac-sha2-512",
		"hmac-sha2-256",
		"hmac-sha1",
		"hmac-sha1-96",
		"hmac-md5",
	}, ",")
	comp := "none,zlib@openssh.com,zlib"

	w := newPayloadWriter()
	w.writeByte(msgKexInit)
	w.writeBytes(cookie[:])
	w.writeString(kex)
	w.writeString(hostKey)
	w.writeString(ciphers)
	w.writeString(ciphers)
	w.writeString(macs)
	w.writeString(macs)
	w.writeString(comp)
	w.writeString(comp)
	w.writeString("") // languages c2s
	w.writeString("") // languages s2c
	w.writeByte(0)    // first_kex_packet_follows
	w.writeUint32(0)  // reserved
	return w.bytes()
}

func mustRand(b []byte) error {
	_, err := rand.Read(b)
	return err
}

// parseKexInit parses a server KEXINIT payload. Validation is minimal:
// we do not reject over-long algorithm lists, we just trim them at the
// RFC ceiling so a hostile server can't make us allocate unbounded
// amounts of memory.
func parseKexInit(payload []byte) (*kexInitPayload, error) {
	if len(payload) < 1 || payload[0] != msgKexInit {
		return nil, fmt.Errorf("not a KEXINIT packet (first byte = %d)", func() byte {
			if len(payload) == 0 {
				return 0
			}
			return payload[0]
		}())
	}
	r := newPayloadReader(payload[1:])
	out := &kexInitPayload{}
	if err := r.readBytes(out.Cookie[:]); err != nil {
		return nil, err
	}
	var err error
	if out.KexAlgorithms, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.ServerHostKeyAlgorithms, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.EncryptionAlgorithmsClientToSvr, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.EncryptionAlgorithmsSvrToClient, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.MACAlgorithmsClientToSvr, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.MACAlgorithmsSvrToClient, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.CompressionAlgorithmsClientToSv, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.CompressionAlgorithmsSvrToClt, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.LanguagesClientToSvr, err = r.readNameList(); err != nil {
		return nil, err
	}
	if out.LanguagesSvrToClient, err = r.readNameList(); err != nil {
		return nil, err
	}
	b, err := r.readByte()
	if err != nil {
		return nil, err
	}
	out.FirstKexPacketFollows = b != 0
	return out, nil
}

// payloadWriter/Reader are tiny helpers for SSH wire encoding. We only
// ever use uint32-prefixed strings and comma-separated name-lists.

type payloadWriter struct{ buf []byte }

func newPayloadWriter() *payloadWriter { return &payloadWriter{} }
func (w *payloadWriter) bytes() []byte { return w.buf }

func (w *payloadWriter) writeByte(b byte)    { w.buf = append(w.buf, b) }
func (w *payloadWriter) writeBytes(b []byte) { w.buf = append(w.buf, b...) }
func (w *payloadWriter) writeUint32(v uint32) {
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], v)
	w.buf = append(w.buf, b[:]...)
}
func (w *payloadWriter) writeString(s string) {
	w.writeUint32(uint32(len(s)))
	w.buf = append(w.buf, s...)
}

type payloadReader struct {
	buf []byte
	pos int
}

func newPayloadReader(b []byte) *payloadReader { return &payloadReader{buf: b} }

func (r *payloadReader) readByte() (byte, error) {
	if r.pos >= len(r.buf) {
		return 0, io.ErrUnexpectedEOF
	}
	b := r.buf[r.pos]
	r.pos++
	return b, nil
}

func (r *payloadReader) readBytes(dst []byte) error {
	if r.pos+len(dst) > len(r.buf) {
		return io.ErrUnexpectedEOF
	}
	copy(dst, r.buf[r.pos:r.pos+len(dst)])
	r.pos += len(dst)
	return nil
}

func (r *payloadReader) readUint32() (uint32, error) {
	if r.pos+4 > len(r.buf) {
		return 0, io.ErrUnexpectedEOF
	}
	v := binary.BigEndian.Uint32(r.buf[r.pos : r.pos+4])
	r.pos += 4
	return v, nil
}

func (r *payloadReader) readNameList() ([]string, error) {
	n, err := r.readUint32()
	if err != nil {
		return nil, err
	}
	if int(n) > len(r.buf)-r.pos {
		return nil, io.ErrUnexpectedEOF
	}
	s := string(r.buf[r.pos : r.pos+int(n)])
	r.pos += int(n)
	if s == "" {
		return nil, nil
	}
	return strings.Split(s, ","), nil
}