package main

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math/big"
	"net"
	"strconv"
	"time"
)

// The network ("listen") transport is an alternative to ssh for reaching an
// endpoint that has no ssh server — typically a Windows source started with
// `clonetool listen`. It is:
//
//   - TLS 1.3 for confidentiality. The listener uses a fresh, self-signed,
//     in-memory certificate generated at startup; the dialer does not verify
//     it (there is no CA). Authentication is by shared password, below.
//   - A password challenge-response *bound to the TLS session*. Both sides
//     derive a key from the password with PBKDF2 and prove knowledge of it by
//     HMAC'ing the connection's RFC 5705 exported keying material, which is
//     unique to this specific TLS session. A man-in-the-middle terminating TLS
//     sees a different exporter value on each leg, so it cannot produce a valid
//     HMAC for either side even though the certificate is unverified. Both
//     directions are checked, so each end authenticates the other.
//
// Once authenticated, the connection carries the ordinary frame protocol
// (frame.go): first a msgRole request naming the endpoint to run, then that
// role's traffic.

const (
	authExporterLabel = "EXPORTER-clonetool-auth-v1"
	authKDFIterations = 200_000
	authKeyLen        = 32
	authExporterLen   = 32
	authMsgType       = "auth"
)

// authKDFSalt is a fixed application salt. Per-connection uniqueness comes from
// the TLS exporter that the HMAC is computed over, not from this salt; it only
// domain-separates the derived key.
var authKDFSalt = []byte("clonetool-listen-auth-salt-v1")

func deriveAuthKey(password string) ([]byte, error) {
	return pbkdf2.Key(sha256.New, password, authKDFSalt, authKDFIterations, authKeyLen)
}

// authTag is HMAC(key, side || exporter): "side" separates the client's proof
// from the server's so the same value can't simply be reflected.
func authTag(key []byte, side string, exporter []byte) []byte {
	m := hmac.New(sha256.New, key)
	m.Write([]byte(side))
	m.Write(exporter)
	return m.Sum(nil)
}

func exporterFor(cs tls.ConnectionState) ([]byte, error) {
	return cs.ExportKeyingMaterial(authExporterLabel, nil, authExporterLen)
}

// newServerTLSConfig builds a TLS config with a fresh self-signed cert.
func newServerTLSConfig() (*tls.Config, error) {
	cert, err := selfSignedCert()
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates: []tls.Certificate{cert},
		MinVersion:   tls.VersionTLS13,
	}, nil
}

func selfSignedCert() (tls.Certificate, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return tls.Certificate{}, err
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return tls.Certificate{}, err
	}
	tmpl := x509.Certificate{
		SerialNumber:          serial,
		Subject:               pkix.Name{CommonName: "clonetool"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(24 * 365 * time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &key.PublicKey, key)
	if err != nil {
		return tls.Certificate{}, err
	}
	return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, nil
}

// dialAuth dials addr, completes the TLS handshake and the mutual password
// challenge, and returns the authenticated connection plus its frame
// reader/writer. On any failure it closes the connection.
func dialAuth(addr, password string, timeout time.Duration) (*tls.Conn, *FrameReader, *FrameWriter, error) {
	d := net.Dialer{Timeout: timeout}
	raw, err := d.Dial("tcp", addr)
	if err != nil {
		return nil, nil, nil, err
	}
	conn := tls.Client(raw, &tls.Config{InsecureSkipVerify: true, MinVersion: tls.VersionTLS13})
	_ = conn.SetDeadline(time.Now().Add(timeout))
	if err := conn.Handshake(); err != nil {
		_ = conn.Close()
		return nil, nil, nil, fmt.Errorf("tls handshake: %w", err)
	}
	key, err := deriveAuthKey(password)
	if err != nil {
		_ = conn.Close()
		return nil, nil, nil, err
	}
	exporter, err := exporterFor(conn.ConnectionState())
	if err != nil {
		_ = conn.Close()
		return nil, nil, nil, fmt.Errorf("tls keying material: %w", err)
	}
	fw := NewFrameWriter(conn)
	fr := NewFrameReader(conn)

	// Prove ourselves to the server, then verify the server back.
	if err := fw.WriteJSON(CtrlMsg{Type: authMsgType, Message: hex.EncodeToString(authTag(key, "client", exporter))}); err != nil {
		_ = conn.Close()
		return nil, nil, nil, err
	}
	m, err := readCtrlFrame(fr)
	if err != nil {
		_ = conn.Close()
		return nil, nil, nil, fmt.Errorf("authentication failed (bad password?): %w", err)
	}
	want := authTag(key, "server", exporter)
	got, _ := hex.DecodeString(m.Message)
	if m.Type != authMsgType || !hmac.Equal(got, want) {
		_ = conn.Close()
		return nil, nil, nil, fmt.Errorf("server authentication failed (password mismatch)")
	}
	_ = conn.SetDeadline(time.Time{}) // clear; per-op timeouts are handled by callers
	return conn, fr, fw, nil
}

// serverAuth runs the server side of the password challenge on an
// already-TLS-handshaken connection.
func serverAuth(conn *tls.Conn, password string, timeout time.Duration) (*FrameReader, *FrameWriter, error) {
	key, err := deriveAuthKey(password)
	if err != nil {
		return nil, nil, err
	}
	exporter, err := exporterFor(conn.ConnectionState())
	if err != nil {
		return nil, nil, fmt.Errorf("tls keying material: %w", err)
	}
	fw := NewFrameWriter(conn)
	fr := NewFrameReader(conn)

	_ = conn.SetReadDeadline(time.Now().Add(timeout))
	m, err := readCtrlFrame(fr)
	if err != nil {
		return nil, nil, err
	}
	want := authTag(key, "client", exporter)
	got, _ := hex.DecodeString(m.Message)
	if m.Type != authMsgType || !hmac.Equal(got, want) {
		return nil, nil, fmt.Errorf("client authentication failed (password mismatch)")
	}
	if err := fw.WriteJSON(CtrlMsg{Type: authMsgType, Message: hex.EncodeToString(authTag(key, "server", exporter))}); err != nil {
		return nil, nil, err
	}
	_ = conn.SetReadDeadline(time.Time{})
	return fr, fw, nil
}

// readCtrlFrame reads one frameCtrlJSON frame and decodes it into a CtrlMsg.
func readCtrlFrame(fr *FrameReader) (CtrlMsg, error) {
	typ, payload, err := fr.ReadFrame()
	if err != nil {
		return CtrlMsg{}, err
	}
	if typ != frameCtrlJSON {
		return CtrlMsg{}, fmt.Errorf("expected control frame, got type %d", typ)
	}
	var m CtrlMsg
	if err := json.Unmarshal(payload, &m); err != nil {
		return CtrlMsg{}, err
	}
	return m, nil
}

// dialPeerLink connects to a listening clonetool over TLS, authenticates,
// requests role for the given window, and waits for its READY handshake. It is
// the net counterpart of spawnPeerLink (see agent.go).
func dialPeerLink(req CtrlMsg, role string) (*peerLink, string) {
	addr := net.JoinHostPort(req.PeerHost, strconv.Itoa(req.PeerPort))
	timeout := time.Duration(req.ConnectTimeoutSec+2) * time.Second
	conn, fr, fw, err := dialAuth(addr, req.PeerPassword, timeout)
	if err != nil {
		return nil, err.Error()
	}
	if err := fw.WriteJSON(CtrlMsg{
		Type: msgRole, Role: role,
		Path: req.PeerPath, Base: req.PeerBase, Size: req.Size, BlockSize: req.BlockSize,
	}); err != nil {
		_ = conn.Close()
		return nil, err.Error()
	}
	if err := waitPeerReady(fr, timeout); err != nil {
		_ = conn.Close()
		return nil, err.Error()
	}
	return &peerLink{
		fw:        fw,
		fr:        fr,
		stderr:    func() string { return "" },
		finish:    conn.Close,
		kill:      func() { _ = conn.Close() },
		closeSend: conn.Close,
	}, ""
}

// waitPeerReady waits for a data role's frameReady, surfacing a msgRoleErr the
// listener may send instead when it can't start the requested role.
func waitPeerReady(fr *FrameReader, timeout time.Duration) error {
	type result struct {
		typ     frameType
		payload []byte
		err     error
	}
	ch := make(chan result, 1)
	go func() {
		typ, payload, err := fr.ReadFrame()
		ch <- result{typ, payload, err}
	}()
	select {
	case r := <-ch:
		if r.err != nil {
			return fmt.Errorf("handshake failed: %w", r.err)
		}
		switch r.typ {
		case frameReady:
			return nil
		case frameCtrlJSON:
			var m CtrlMsg
			if json.Unmarshal(r.payload, &m) == nil && m.Type == msgRoleErr {
				return fmt.Errorf("remote refused role: %s", m.Message)
			}
			return fmt.Errorf("handshake failed: unexpected control frame")
		default:
			return fmt.Errorf("handshake failed: unexpected frame type %d", r.typ)
		}
	case <-time.After(timeout):
		return fmt.Errorf("handshake timed out after %s", timeout)
	}
}