package main

import (
	"bufio"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	agentControlLineMax = 16 * 1024
	agentPipeTokenEnv   = "REMOTEDESK_AGENT_PIPE_TOKEN"
	agentPipeAuthDomain = "RemoteDesk Windows Agent pipe auth v1\x00"
)

type mediaRingDescriptor struct {
	SchemaVersion  uint16               `json:"schema_version"`
	SessionID      string               `json:"session_id"`
	Generation     uint64               `json:"generation"`
	OwnerPID       uint32               `json:"owner_pid"`
	RingPath       string               `json:"ring_path"`
	SlotCount      uint32               `json:"slot_count"`
	SlotBytes      uint32               `json:"slot_bytes"`
	VideoTransport string               `json:"video_transport"`
	AudioTransport string               `json:"audio_transport"`
	AudioRing      *audioRingDescriptor `json:"audio_ring,omitempty"`
}

type audioRingDescriptor struct {
	SchemaVersion  uint16 `json:"schema_version"`
	SessionID      string `json:"session_id"`
	Generation     uint64 `json:"generation"`
	OwnerPID       uint32 `json:"owner_pid"`
	RingPath       string `json:"ring_path"`
	SlotCount      uint32 `json:"slot_count"`
	SlotBytes      uint32 `json:"slot_bytes"`
	AudioTransport string `json:"audio_transport"`
}

type encodedDesktopOpened struct {
	Kind           string              `json:"kind"`
	Error          string              `json:"error"`
	VideoCodec     string              `json:"video_codec"`
	MediaTransport string              `json:"media_transport"`
	Ring           mediaRingDescriptor `json:"ring"`
}

type agentHello struct {
	Kind          string `json:"kind"`
	Authenticated bool   `json:"authenticated"`
}

type agentAuthChallenge struct {
	Kind            string `json:"kind"`
	ProtocolVersion uint8  `json:"protocol_version"`
	Nonce           string `json:"nonce"`
}

func openWindowsEncodedSession(address, sessionID string, fps uint8, allowSoftwareFallback bool) (net.Conn, mediaRingDescriptor, error) {
	if err := validateAgentControlAddress(address); err != nil {
		return nil, mediaRingDescriptor{}, err
	}
	if !validSessionID(sessionID) || fps == 0 || fps > 120 {
		return nil, mediaRingDescriptor{}, errors.New("Windows Agent encoded session parameters are invalid")
	}
	isPipe := isWindowsAgentPipe(address)
	connection, err := dialWindowsAgentControl(address, 2*time.Second)
	if err != nil {
		return nil, mediaRingDescriptor{}, fmt.Errorf("connect Windows Agent control: %w", err)
	}
	closeOnError := true
	defer func() {
		if closeOnError {
			_ = connection.Close()
		}
	}()
	_ = connection.SetDeadline(time.Now().Add(10 * time.Second))
	reader := bufio.NewReaderSize(connection, agentControlLineMax+1)
	if isPipe {
		if err := authenticateWindowsAgentPipe(reader, connection); err != nil {
			return nil, mediaRingDescriptor{}, err
		}
	}
	var hello agentHello
	if err := readBoundedAgentJSON(reader, &hello); err != nil || hello.Kind != "windows_agent_hello" || (isPipe && !hello.Authenticated) {
		return nil, mediaRingDescriptor{}, errors.New("Windows Agent control hello is invalid")
	}
	request := struct {
		Kind                  string `json:"kind"`
		SessionID             string `json:"session_id"`
		CaptureMode           string `json:"capture_mode"`
		VideoCodec            string `json:"video_codec"`
		FramesPerSecond       uint8  `json:"frames_per_second"`
		AllowSoftwareFallback bool   `json:"allow_software_fallback"`
	}{
		Kind: "open_encoded_desktop", SessionID: sessionID,
		CaptureMode: "compatibility", VideoCodec: "h264", FramesPerSecond: fps,
		AllowSoftwareFallback: allowSoftwareFallback,
	}
	encoded, err := json.Marshal(request)
	if err != nil {
		return nil, mediaRingDescriptor{}, err
	}
	encoded = append(encoded, '\n')
	if _, err := connection.Write(encoded); err != nil {
		return nil, mediaRingDescriptor{}, fmt.Errorf("open Windows encoded desktop: %w", err)
	}
	var opened encodedDesktopOpened
	if err := readBoundedAgentJSON(reader, &opened); err != nil {
		return nil, mediaRingDescriptor{}, err
	}
	if opened.Kind != "encoded_desktop_opened" || opened.VideoCodec != "h264" ||
		opened.MediaTransport != "shared_memory_encoded_access_units" {
		return nil, mediaRingDescriptor{}, fmt.Errorf("Windows encoded desktop unavailable: %s", opened.Error)
	}
	if err := validateMediaRingDescriptor(opened.Ring, sessionID); err != nil {
		return nil, mediaRingDescriptor{}, err
	}
	_ = connection.SetDeadline(time.Time{})
	closeOnError = false
	return connection, opened.Ring, nil
}

func dialWindowsAgentControl(address string, timeout time.Duration) (net.Conn, error) {
	if isWindowsAgentPipe(address) {
		return dialWindowsAgentPipe(address, timeout)
	}
	return net.DialTimeout("tcp", address, timeout)
}

func authenticateWindowsAgentPipe(reader *bufio.Reader, connection net.Conn) error {
	var challenge agentAuthChallenge
	if err := readBoundedAgentJSON(reader, &challenge); err != nil {
		return err
	}
	if challenge.Kind != "agent_auth_challenge" || challenge.ProtocolVersion != 1 {
		return errors.New("Windows Agent Pipe challenge is invalid")
	}
	nonce, err := base64.RawURLEncoding.DecodeString(challenge.Nonce)
	if err != nil || len(nonce) != 32 {
		return errors.New("Windows Agent Pipe nonce is invalid")
	}
	encodedToken, ok := os.LookupEnv(agentPipeTokenEnv)
	if !ok {
		return fmt.Errorf("%s is required for authenticated Windows Agent Pipe", agentPipeTokenEnv)
	}
	_ = os.Unsetenv(agentPipeTokenEnv)
	key, err := base64.RawURLEncoding.DecodeString(encodedToken)
	if err != nil || len(key) != 32 {
		return errors.New("Windows Agent Pipe token must be 32-byte base64url")
	}
	mac := windowsAgentPipeMAC(key, nonce)
	request := struct {
		Kind            string `json:"kind"`
		ProtocolVersion uint8  `json:"protocol_version"`
		MAC             string `json:"mac"`
	}{
		Kind:            "agent_auth",
		ProtocolVersion: 1,
		MAC:             base64.RawURLEncoding.EncodeToString(mac),
	}
	encoded, err := json.Marshal(request)
	if err != nil {
		return err
	}
	encoded = append(encoded, '\n')
	if _, err := connection.Write(encoded); err != nil {
		return fmt.Errorf("authenticate Windows Agent Pipe: %w", err)
	}
	return nil
}

func windowsAgentPipeMAC(key, nonce []byte) []byte {
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(agentPipeAuthDomain))
	_, _ = mac.Write(nonce)
	return mac.Sum(nil)
}

func readBoundedAgentJSON(reader *bufio.Reader, output any) error {
	line, err := reader.ReadSlice('\n')
	if err != nil {
		return fmt.Errorf("read Windows Agent control: %w", err)
	}
	if len(line) == 0 || len(line) > agentControlLineMax || line[len(line)-1] != '\n' {
		return errors.New("Windows Agent control response exceeds its bound")
	}
	if err := json.Unmarshal(line, output); err != nil {
		return errors.New("Windows Agent control response is invalid JSON")
	}
	return nil
}

func validateAgentControlAddress(address string) error {
	if isWindowsAgentPipe(address) {
		if len(address) > 256 || !strings.HasPrefix(strings.ToLower(address), `\\.\pipe\remotedesk\`) || strings.Contains(address, "..") {
			return errors.New("Windows Agent Pipe name is outside the RemoteDesk namespace")
		}
		return nil
	}
	host, port, err := net.SplitHostPort(address)
	if err != nil || port == "" {
		return errors.New("Windows Agent control must be an IP socket address")
	}
	ip := net.ParseIP(host)
	if ip == nil || !ip.IsLoopback() {
		return errors.New("Windows Agent control must use a loopback IP")
	}
	return nil
}

func isWindowsAgentPipe(address string) bool {
	return strings.HasPrefix(strings.ToLower(address), `\\.\pipe\`)
}

func validateMediaRingDescriptor(descriptor mediaRingDescriptor, sessionID string) error {
	if descriptor.SchemaVersion != 1 || descriptor.SessionID != sessionID || descriptor.Generation == 0 || descriptor.OwnerPID == 0 ||
		descriptor.VideoTransport != "shared_memory_encoded_access_units" || descriptor.AudioTransport != "separate_transport" ||
		!filepath.IsAbs(descriptor.RingPath) || descriptor.SlotCount < 2 || descriptor.SlotCount > maxRingSlots ||
		descriptor.SlotCount&(descriptor.SlotCount-1) != 0 || descriptor.SlotBytes == 0 || descriptor.SlotBytes > maxRingSlotBytes {
		return errors.New("Windows Agent media ring descriptor is invalid")
	}
	if descriptor.AudioRing != nil {
		audio := descriptor.AudioRing
		if audio.SchemaVersion != 1 || audio.SessionID != sessionID || audio.Generation != descriptor.Generation ||
			audio.OwnerPID != descriptor.OwnerPID || audio.AudioTransport != "shared_memory_opus_packets" ||
			!filepath.IsAbs(audio.RingPath) || audio.SlotCount < 2 || audio.SlotCount > maxRingSlots ||
			audio.SlotCount&(audio.SlotCount-1) != 0 || audio.SlotBytes == 0 || audio.SlotBytes > maxRingSlotBytes {
			return errors.New("Windows Agent audio ring descriptor is invalid")
		}
	}
	return nil
}

func validSessionID(value string) bool {
	if len(value) == 0 || len(value) > 128 {
		return false
	}
	for _, character := range []byte(value) {
		if !((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') ||
			(character >= '0' && character <= '9') || character == '-' || character == '_') {
			return false
		}
	}
	return true
}