// Package sshproxy terminates SSH authentication at the gateway and forwards
// permitted session channels over an independently authenticated upstream.
package sshproxy

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"linkport/internal/model"
	"linkport/internal/outbound"
	"linkport/internal/protocol/transport"
)

type Handler struct{}

func New() *Handler { return &Handler{} }

var _ model.Handler = (*Handler)(nil)

func GenerateHostKey() (string, error) {
	_, key, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return "", err
	}
	der, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		return "", err
	}
	return string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})), nil
}

func targetAddress(p model.Proxy) string {
	return net.JoinHostPort(p.TargetHost, strconv.Itoa(p.TargetPort))
}
func supportedKey(k ssh.PublicKey) bool {
	return k.Type() == ssh.KeyAlgoED25519 || k.Type() == ssh.KeyAlgoRSA
}
func equalSecret(a, b string) bool {
	x, y := sha256.Sum256([]byte(a)), sha256.Sum256([]byte(b))
	return subtle.ConstantTimeCompare(x[:], y[:]) == 1
}
func algorithms() ssh.Config {
	return ssh.Config{KeyExchanges: []string{"curve25519-sha256", "ecdh-sha2-nistp256", "ecdh-sha2-nistp384", "ecdh-sha2-nistp521", "diffie-hellman-group14-sha256"}, Ciphers: []string{"chacha20-poly1305@openssh.com", "aes128-gcm@openssh.com", "aes256-gcm@openssh.com", "aes128-ctr", "aes192-ctr", "aes256-ctr"}, MACs: []string{"hmac-sha2-256-etm@openssh.com", "hmac-sha2-512-etm@openssh.com", "hmac-sha2-256", "hmac-sha2-512"}}
}
func serverConfig(p model.Proxy) (*ssh.ServerConfig, error) {
	signer, err := ssh.ParsePrivateKey([]byte(p.SSH.HostKey))
	if err != nil {
		return nil, errors.New("invalid gateway host key")
	}
	if signer.PublicKey().Type() != ssh.KeyAlgoED25519 {
		return nil, errors.New("gateway host key must be Ed25519")
	}
	cfg := &ssh.ServerConfig{Config: algorithms(), MaxAuthTries: 5, ServerVersion: "SSH-2.0-linkport", PublicKeyAuthAlgorithms: []string{ssh.KeyAlgoED25519, ssh.KeyAlgoRSASHA256, ssh.KeyAlgoRSASHA512}}
	cfg.AddHostKey(signer)
	if p.GatewayPassword != "" || len(p.GatewayAccounts) > 0 {
		cfg.PasswordCallback = func(c ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			a, found := transport.Account(p, c.User(), c.RemoteAddr())
			if found && a.Password != "" && equalSecret(string(password), a.Password) {
				return nil, nil
			}
			return nil, errors.New("authentication failed")
		}
	}
	entries := p.SSH.AuthorizedKeys
	for _, a := range p.GatewayAccounts {
		entries = append(append([]string(nil), entries...), a.AuthorizedKeys...)
	}
	for _, entry := range entries {
		key, _, options, rest, e := ssh.ParseAuthorizedKey([]byte(entry))
		if e != nil || len(options) > 0 || len(bytes.TrimSpace(rest)) > 0 || !supportedKey(key) {
			return nil, errors.New("invalid or unsupported authorized key")
		}
	}
	if len(entries) > 0 {
		cfg.PublicKeyCallback = func(c ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			a, found := transport.Account(p, c.User(), c.RemoteAddr())
			if found {
				for _, entry := range a.AuthorizedKeys {
					known, _, _, _, err := ssh.ParseAuthorizedKey([]byte(entry))
					if err == nil && bytes.Equal(known.Marshal(), key.Marshal()) {
						return nil, nil
					}
				}
			}
			return nil, errors.New("authentication failed")
		}
	}
	if cfg.PasswordCallback == nil && cfg.PublicKeyCallback == nil {
		return nil, errors.New("gateway authentication is not configured")
	}
	return cfg, nil
}
func upstreamConfig(p model.Proxy) (*ssh.ClientConfig, error) {
	if err := p.CheckTargetCredential(); err != nil {
		return nil, err
	}
	var callback ssh.HostKeyCallback
	var err error
	if p.SSH.IgnoreHostKey {
		callback = ssh.InsecureIgnoreHostKey()
	} else {
		key, _, _, rest, parseErr := ssh.ParseAuthorizedKey([]byte(p.SSH.TargetHostKey))
		err = parseErr
		if err != nil || len(bytes.TrimSpace(rest)) > 0 || !supportedKey(key) {
			return nil, model.AtConnectionStage("ssh_host_key_unconfirmed", errors.New("target host key must be confirmed"))
		}
		verify := ssh.FixedHostKey(key)
		callback = func(host string, addr net.Addr, observed ssh.PublicKey) error {
			if err := verify(host, addr, observed); err != nil {
				return model.AtConnectionStage("ssh_host_identity", err)
			}
			return nil
		}
	}
	cfg := &ssh.ClientConfig{Config: algorithms(), User: p.TargetUser, HostKeyAlgorithms: []string{ssh.KeyAlgoED25519, ssh.KeyAlgoRSASHA512, ssh.KeyAlgoRSASHA256}, HostKeyCallback: callback}
	switch p.SSH.AuthMethod {
	case "password":
		cfg.Auth = []ssh.AuthMethod{ssh.Password(p.TargetPassword)}
	case "privatekey", "privateKey":
		var signer ssh.Signer
		if p.SSH.Passphrase != "" {
			signer, err = ssh.ParsePrivateKeyWithPassphrase([]byte(p.SSH.PrivateKey), []byte(p.SSH.Passphrase))
		} else {
			signer, err = ssh.ParsePrivateKey([]byte(p.SSH.PrivateKey))
		}
		if err != nil {
			return nil, errors.New("invalid target private key or passphrase")
		}
		if !supportedKey(signer.PublicKey()) {
			return nil, errors.New("unsupported target private key")
		}
		if signer.PublicKey().Type() == ssh.KeyAlgoRSA {
			multi, ok := signer.(ssh.AlgorithmSigner)
			if !ok {
				return nil, errors.New("RSA SHA-2 signing unavailable")
			}
			signer, err = ssh.NewSignerWithAlgorithms(multi, []string{ssh.KeyAlgoRSASHA512, ssh.KeyAlgoRSASHA256})
			if err != nil {
				return nil, err
			}
		}
		signer, err = model.SSHCertificateSigner(p.SSH.Certificate, signer)
		if err != nil {
			return nil, errors.New("invalid target SSH certificate or private key")
		}
		cfg.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	default:
		return nil, errors.New("invalid target authentication method")
	}
	return cfg, nil
}

// connect leaves the connection owned by the caller; cancellation remains active
// throughout the session, not only while dialing.
func connect(ctx context.Context, p model.Proxy, dial model.DialFunc, cfg *ssh.ClientConfig) (*ssh.Client, func(), error) {
	raw, err := dial(ctx, "tcp", targetAddress(p))
	if err != nil {
		return nil, nil, fmt.Errorf("SSH target connection: %w", err)
	}
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	cleanup := func() { stop(); raw.Close() }
	raw.SetDeadline(time.Now().Add(15 * time.Second))
	conn, channels, requests, err := ssh.NewClientConn(raw, targetAddress(p), cfg)
	if err != nil {
		cleanup()
		return nil, nil, fmt.Errorf("SSH target handshake: %w", err)
	}
	raw.SetDeadline(time.Time{})
	return ssh.NewClient(conn, channels, requests), cleanup, nil
}

func (h *Handler) Test(ctx context.Context, p model.Proxy, dial model.DialFunc) (model.TestResult, error) {
	cfg, err := upstreamConfig(p)
	if err != nil {
		var failure *model.ConnectionFailure
		if errors.As(err, &failure) && failure.Stage == "ssh_host_key_unconfirmed" {
			outbound.Begin(ctx, "ssh_host_key_unconfirmed")(err)
		}
		return model.TestResult{}, err
	}
	var fingerprint string
	verify := cfg.HostKeyCallback
	cfg.HostKeyCallback = func(host string, addr net.Addr, key ssh.PublicKey) error {
		done := outbound.Begin(ctx, "ssh_host_identity")
		if err := verify(host, addr, key); err != nil {
			done(err)
			return err
		}
		done(nil)
		fingerprint = ssh.FingerprintSHA256(key)
		return nil
	}
	c, closeConn, err := connect(ctx, p, dial, cfg)
	if err != nil {
		return model.TestResult{}, err
	}
	defer closeConn()
	defer c.Close()
	if p.Type == "sftp" {
		files, err := sftp.NewClient(c)
		if err != nil {
			return model.TestResult{}, errors.New("SFTP subsystem unavailable")
		}
		defer files.Close()
		if _, err := files.Getwd(); err != nil {
			return model.TestResult{}, errors.New("SFTP working directory unavailable")
		}
	}
	result := model.TestResult{Version: string(c.ServerVersion()), Fingerprint: fingerprint}
	return result, nil
}

// Probe deliberately aborts in HostKeyCallback, before user authentication.
func Probe(ctx context.Context, p model.Proxy, dial model.DialFunc) (model.TestResult, error) {
	var result model.TestResult
	captured := errors.New("host key captured")
	cfg := &ssh.ClientConfig{Config: algorithms(), User: "linkport-host-key-probe", HostKeyAlgorithms: []string{ssh.KeyAlgoED25519, ssh.KeyAlgoRSASHA512, ssh.KeyAlgoRSASHA256}, HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
		result.Fingerprint = ssh.FingerprintSHA256(key)
		result.HostKey = string(ssh.MarshalAuthorizedKey(key))
		return captured
	}}
	c, cleanup, err := connect(ctx, p, dial, cfg)
	if cleanup != nil {
		defer cleanup()
	}
	if c != nil {
		defer c.Close()
	}
	if errors.Is(err, captured) {
		return result, nil
	}
	return model.TestResult{}, err
}

func (h *Handler) Serve(ctx context.Context, raw net.Conn, p model.Proxy, dial model.DialFunc, _ *tls.Config, event model.EventFunc) error {
	defer raw.Close()
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	defer stop()
	cfg, err := serverConfig(p)
	if err != nil {
		return err
	}
	cfg.AuthLogCallback = func(c ssh.ConnMetadata, method string, err error) {
		if method != "none" {
			transport.AuthenticationEvent(event, c.User(), err == nil)
		}
	}
	raw.SetDeadline(time.Now().Add(15 * time.Second))
	down, channels, requests, err := ssh.NewServerConn(raw, cfg)
	if err != nil {
		return errors.New("SSH gateway authentication or handshake failed")
	}
	defer down.Close()
	account, found := transport.Account(p, down.User(), raw.RemoteAddr())
	if !found {
		return errors.New("SSH gateway account unavailable")
	}
	p = account.Apply(p)
	dial = p.DialForTarget(dial)
	upstreamCfg, err := upstreamConfig(p)
	if err != nil {
		return err
	}
	raw.SetDeadline(time.Time{})
	sessionCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	upstream, cleanup, err := connect(sessionCtx, p, dial, upstreamCfg)
	if err != nil {
		return err
	}
	defer cleanup()
	defer upstream.Close()
	go func() { upstream.Wait(); down.Close() }()
	go func() {
		for req := range requests {
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}()
	slots := make(chan struct{}, 10)
	var sessions sync.WaitGroup
	defer func() { upstream.Close(); down.Close(); sessions.Wait() }()
	for incoming := range channels {
		if incoming.ChannelType() != "session" {
			incoming.Reject(ssh.Prohibited, "only session channels are supported")
			if event != nil {
				event(model.EventProtocolRejected, "ssh:channel_not_allowed")
			}
			continue
		}
		select {
		case slots <- struct{}{}:
		default:
			incoming.Reject(ssh.ResourceShortage, "session limit reached")
			continue
		}
		sessions.Add(1)
		go func(ch ssh.NewChannel) {
			defer sessions.Done()
			defer func() { <-slots }()
			bridgeSession(ctx, upstream, ch, event, p.Type == "sftp")
		}(incoming)
	}
	return nil
}

func allowedRequest(req *ssh.Request) bool {
	switch req.Type {
	case "pty-req", "window-change", "shell", "exec", "signal", "break":
		return true
	case "subsystem":
		var v struct{ Name string }
		return ssh.Unmarshal(req.Payload, &v) == nil && v.Name == "sftp"
	case "env":
		var v struct{ Name, Value string }
		return ssh.Unmarshal(req.Payload, &v) == nil && (v.Name == "TERM" || v.Name == "LANG" || strings.HasPrefix(v.Name, "LC_"))
	default:
		return false
	}
}
func bridgeSession(ctx context.Context, up *ssh.Client, incoming ssh.NewChannel, event model.EventFunc, sftpOnly bool) {
	target, targetRequests, err := up.OpenChannel("session", incoming.ExtraData())
	if err != nil {
		incoming.Reject(ssh.ConnectionFailed, "target session unavailable")
		return
	}
	defer target.Close()
	client, clientRequests, err := incoming.Accept()
	if err != nil {
		return
	}
	defer client.Close()
	input := newInputAudit(ctx, client)
	if sftpOnly {
		input.mode.Store(2)
	}
	var workers sync.WaitGroup
	var requestMu sync.Mutex
	var output sync.WaitGroup
	output.Add(2)
	workers.Add(5)
	go func() {
		defer workers.Done()
		defer target.Close()
		started := false
		for req := range clientRequests {
			requestMu.Lock()
			ok := false
			allowed := allowedRequest(req)
			startup := req.Type == "exec" || req.Type == "shell" || req.Type == "subsystem"
			if startup && started {
				allowed = false
			}
			if sftpOnly {
				allowed = allowed && req.Type == "subsystem"
			}
			decision := "blocked"
			if allowed {
				decision = "allowed"
			}
			if err := auditRequest(ctx, req, decision, input.output); err != nil {
				input.output.fail()
				client.Close()
				requestMu.Unlock()
				return
			}
			if allowed {
				if startup {
					started = true
				}
				if req.Type == "shell" {
					input.mode.Store(1)
				}
				if req.Type == "subsystem" {
					input.mode.Store(2)
				}
				var sendErr error
				ok, sendErr = target.SendRequest(req.Type, req.WantReply, req.Payload)
				if startup && (sendErr != nil || (req.WantReply && !ok)) {
					input.mode.Store(0)
					state := "rejected"
					if sendErr != nil {
						state = "interrupted"
					}
					if err := input.output.finish(state); err != nil {
						client.Close()
						target.Close()
					}
				}
				if req.Type == "pty-req" && ok {
					input.pty.Store(true)
				}
			} else if event != nil {
				event(model.EventProtocolRejected, "ssh:request_not_allowed")
			}
			if req.WantReply {
				req.Reply(ok, nil)
			}
			requestMu.Unlock()
		}
	}()
	go func() {
		defer workers.Done()
		if _, err := io.Copy(target, input); err != nil {
			input.output.fail()
			target.Close()
			client.Close()
		}
		target.CloseWrite()
	}()
	go func() {
		defer workers.Done()
		defer output.Done()
		if _, err := io.Copy(client.Stderr(), outputAudit{input: input, reader: target.Stderr(), stream: "stderr"}); err != nil {
			input.output.fail()
			target.Close()
			client.Close()
		}
	}()
	go func() {
		defer workers.Done()
		defer output.Done()
		if _, err := io.Copy(client, outputAudit{input: input, reader: target, stream: "stdout"}); err != nil {
			input.output.fail()
			target.Close()
			client.Close()
		}
	}()
	go func() { defer workers.Done(); output.Wait(); client.CloseWrite() }()
	// Keep request processing alive until upstream closes its channel so exit-status
	// cannot be lost when stdout reaches EOF before the exit notification.
	for req := range targetRequests {
		input.output.exit(req)
		ok := false
		if req.Type == "exit-status" || req.Type == "exit-signal" || req.Type == "xon-xoff" {
			ok, _ = client.SendRequest(req.Type, req.WantReply, req.Payload)
		}
		if req.WantReply {
			req.Reply(ok, nil)
		}
	}
	target.Close()
	output.Wait()
	// Target channel closure may precede forwarding the reply to exec/subsystem.
	// Wait for that in-flight reply before closing the downstream channel.
	requestMu.Lock()
	client.Close()
	requestMu.Unlock()
	workers.Wait()
	if err := input.output.finish(""); err != nil && event != nil {
		event(model.EventProtocolRejected, "ssh:output_audit_unavailable")
	}
}