package externalauth

import (
	"context"
	"crypto/md5"
	"crypto/rand"
	"crypto/subtle"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"time"
)

type Server struct {
	Host   string
	Port   int
	Secret string
}

type Authenticator interface {
	Authenticate(context.Context, string, Server, string, string, time.Duration) (bool, error)
}

type Client struct{ Dialer net.Dialer }

func (c Client) Authenticate(ctx context.Context, method string, server Server, username, password string, timeout time.Duration) (bool, error) {
	if timeout <= 0 {
		timeout = 5 * time.Second
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	switch strings.ToLower(strings.TrimSpace(method)) {
	case "radius":
		return c.radius(ctx, server, username, password)
	case "tacacs", "tacacs+":
		return c.tacacs(ctx, server, username, password)
	default:
		return false, fmt.Errorf("unsupported external authentication method %q", method)
	}
}

func (c Client) radius(ctx context.Context, server Server, username, password string) (bool, error) {
	if err := validateInput(server, username, password); err != nil {
		return false, err
	}
	address := net.JoinHostPort(server.Host, strconv.Itoa(server.Port))
	conn, err := c.Dialer.DialContext(ctx, "udp", address)
	if err != nil {
		return false, fmt.Errorf("dial RADIUS server: %w", err)
	}
	defer conn.Close()
	if deadline, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}
	var requestAuth [16]byte
	if _, err := rand.Read(requestAuth[:]); err != nil {
		return false, fmt.Errorf("generate RADIUS authenticator: %w", err)
	}
	var identifier [1]byte
	if _, err := rand.Read(identifier[:]); err != nil {
		return false, fmt.Errorf("generate RADIUS identifier: %w", err)
	}
	attrs := append(radiusAttribute(1, []byte(username)), radiusAttribute(2, hideRadiusPassword([]byte(server.Secret), requestAuth[:], []byte(password)))...)
	attrs = append(attrs, radiusAttribute(32, []byte("tera-console-server"))...)
	packet := make([]byte, 20+len(attrs))
	packet[0], packet[1] = 1, identifier[0]
	binary.BigEndian.PutUint16(packet[2:4], uint16(len(packet)))
	copy(packet[4:20], requestAuth[:])
	copy(packet[20:], attrs)
	if _, err := conn.Write(packet); err != nil {
		return false, fmt.Errorf("write RADIUS request: %w", err)
	}
	response := make([]byte, 4096)
	n, err := conn.Read(response)
	if err != nil {
		return false, fmt.Errorf("read RADIUS response: %w", err)
	}
	response = response[:n]
	if len(response) < 20 || response[1] != identifier[0] {
		return false, errors.New("invalid RADIUS response header")
	}
	length := int(binary.BigEndian.Uint16(response[2:4]))
	if length != len(response) {
		return false, errors.New("invalid RADIUS response length")
	}
	authInput := append([]byte(nil), response...)
	copy(authInput[4:20], requestAuth[:])
	authInput = append(authInput, []byte(server.Secret)...)
	digest := md5.Sum(authInput)
	if subtle.ConstantTimeCompare(digest[:], response[4:20]) != 1 {
		return false, errors.New("invalid RADIUS response authenticator")
	}
	switch response[0] {
	case 2:
		return true, nil
	case 3:
		return false, nil
	case 11:
		return false, errors.New("RADIUS Access-Challenge is not supported")
	default:
		return false, fmt.Errorf("unexpected RADIUS response code %d", response[0])
	}
}

func radiusAttribute(kind byte, value []byte) []byte {
	result := make([]byte, 2+len(value))
	result[0], result[1] = kind, byte(len(result))
	copy(result[2:], value)
	return result
}

func hideRadiusPassword(secret, requestAuth, password []byte) []byte {
	length := ((len(password) + 15) / 16) * 16
	if length == 0 {
		length = 16
	}
	padded := make([]byte, length)
	copy(padded, password)
	result := make([]byte, length)
	previous := requestAuth
	for offset := 0; offset < length; offset += 16 {
		input := append(append([]byte(nil), secret...), previous...)
		digest := md5.Sum(input)
		for index := 0; index < 16; index++ {
			result[offset+index] = padded[offset+index] ^ digest[index]
		}
		previous = result[offset : offset+16]
	}
	return result
}

func (c Client) tacacs(ctx context.Context, server Server, username, password string) (bool, error) {
	if err := validateInput(server, username, password); err != nil {
		return false, err
	}
	address := net.JoinHostPort(server.Host, strconv.Itoa(server.Port))
	conn, err := c.Dialer.DialContext(ctx, "tcp", address)
	if err != nil {
		return false, fmt.Errorf("dial TACACS+ server: %w", err)
	}
	defer conn.Close()
	if deadline, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}
	var sessionBytes [4]byte
	if _, err := rand.Read(sessionBytes[:]); err != nil {
		return false, err
	}
	sessionID := binary.BigEndian.Uint32(sessionBytes[:])
	start := []byte{1, 1, 1, 1, byte(len(username)), 0, 0, 0}
	start = append(start, []byte(username)...)
	if err := writeTacacsPacket(conn, sessionID, 1, []byte(server.Secret), start); err != nil {
		return false, err
	}
	sequence := byte(2)
	for attempts := 0; attempts < 6; attempts++ {
		body, err := readTacacsPacket(conn, sessionID, sequence, []byte(server.Secret))
		if err != nil {
			return false, err
		}
		if len(body) < 6 {
			return false, errors.New("invalid TACACS+ authentication reply")
		}
		status := body[0]
		switch status {
		case 1:
			return true, nil
		case 2:
			return false, nil
		case 4, 5, 3:
			value := password
			if status == 4 {
				value = username
			}
			continueBody := make([]byte, 5+len(value))
			binary.BigEndian.PutUint16(continueBody[0:2], uint16(len(value)))
			copy(continueBody[5:], []byte(value))
			sequence++
			if err := writeTacacsPacket(conn, sessionID, sequence, []byte(server.Secret), continueBody); err != nil {
				return false, err
			}
			sequence++
		case 6:
			return false, errors.New("TACACS+ server requested authentication restart")
		default:
			return false, fmt.Errorf("TACACS+ authentication error status %d", status)
		}
	}
	return false, errors.New("TACACS+ authentication exchange limit exceeded")
}

func writeTacacsPacket(writer io.Writer, sessionID uint32, sequence byte, secret, body []byte) error {
	header := make([]byte, 12)
	header[0], header[1], header[2] = 0xc0, 1, sequence
	binary.BigEndian.PutUint32(header[4:8], sessionID)
	binary.BigEndian.PutUint32(header[8:12], uint32(len(body)))
	encoded := tacacsXOR(header, secret, body)
	if _, err := writer.Write(append(header, encoded...)); err != nil {
		return fmt.Errorf("write TACACS+ packet: %w", err)
	}
	return nil
}

func readTacacsPacket(reader io.Reader, sessionID uint32, sequence byte, secret []byte) ([]byte, error) {
	header := make([]byte, 12)
	if _, err := io.ReadFull(reader, header); err != nil {
		return nil, fmt.Errorf("read TACACS+ header: %w", err)
	}
	if header[0] != 0xc0 || header[1] != 1 || header[2] != sequence || binary.BigEndian.Uint32(header[4:8]) != sessionID || header[3]&1 != 0 {
		return nil, errors.New("invalid TACACS+ response header")
	}
	length := binary.BigEndian.Uint32(header[8:12])
	if length > 64*1024 {
		return nil, errors.New("TACACS+ response exceeds limit")
	}
	body := make([]byte, length)
	if _, err := io.ReadFull(reader, body); err != nil {
		return nil, fmt.Errorf("read TACACS+ body: %w", err)
	}
	return tacacsXOR(header, secret, body), nil
}

func tacacsXOR(header, secret, body []byte) []byte {
	result := make([]byte, len(body))
	pad := make([]byte, 0, len(body))
	var previous []byte
	for len(pad) < len(body) {
		input := append([]byte(nil), header[4:8]...)
		input = append(input, secret...)
		input = append(input, header[0], header[2])
		input = append(input, previous...)
		digest := md5.Sum(input)
		previous = digest[:]
		pad = append(pad, previous...)
	}
	for index := range body {
		result[index] = body[index] ^ pad[index]
	}
	return result
}

func validateInput(server Server, username, password string) error {
	if strings.TrimSpace(server.Host) == "" || server.Port < 1 || server.Port > 65535 {
		return errors.New("external authentication server address is invalid")
	}
	if server.Secret == "" {
		return errors.New("external authentication shared secret is required")
	}
	if username == "" || len(username) > 253 || password == "" || len(password) > 253 {
		return errors.New("external authentication credentials are invalid")
	}
	return nil
}