// setec-agent — E2E encrypted command executor for SETEC LABS Manager.
//
// Reads AES-256-GCM tunnel key from /etc/setec/tunnel.key (32 bytes, hex-encoded).
// Receives base64-encoded encrypted commands on stdin.
// Decrypts, executes via /bin/sh, encrypts the response, outputs base64.
//
// Protocol:
//   Input:  base64( nonce[12] + AES-GCM-ciphertext )
//   Output: base64( nonce[12] + AES-GCM-ciphertext )
//   Plaintext response is JSON: {"stdout":"...","stderr":"...","exit_code":N}
//
// Build: GOOS=linux GOARCH=amd64 go build -o setec-agent -ldflags="-s -w" .

package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

const keyPath = "/etc/setec/tunnel.key"

type Response struct {
	Stdout   string `json:"stdout"`
	Stderr   string `json:"stderr"`
	ExitCode int    `json:"exit_code"`
}

func loadKey() ([]byte, error) {
	data, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot read tunnel key: %w", err)
	}
	hexKey := strings.TrimSpace(string(data))
	key, err := hex.DecodeString(hexKey)
	if err != nil {
		return nil, fmt.Errorf("invalid tunnel key hex: %w", err)
	}
	if len(key) != 32 {
		return nil, fmt.Errorf("tunnel key must be 32 bytes, got %d", len(key))
	}
	return key, nil
}

func decrypt(key []byte, b64input string) ([]byte, error) {
	raw, err := base64.StdEncoding.DecodeString(strings.TrimSpace(b64input))
	if err != nil {
		return nil, fmt.Errorf("base64 decode: %w", err)
	}
	if len(raw) < 12+16 {
		return nil, fmt.Errorf("ciphertext too short")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := raw[:12]
	ciphertext := raw[12:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decryption failed (bad key or tampered data): %w", err)
	}
	return plaintext, nil
}

func encrypt(key []byte, plaintext []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, 12)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func executeCommand(cmd string) Response {
	ctx_timeout := 120 * time.Second
	command := exec.Command("/bin/sh", "-c", cmd)

	var stdout, stderr strings.Builder
	command.Stdout = &stdout
	command.Stderr = &stderr

	// Start with timeout
	done := make(chan error, 1)
	command.Start()
	go func() { done <- command.Wait() }()

	select {
	case err := <-done:
		exitCode := 0
		if err != nil {
			if exitErr, ok := err.(*exec.ExitError); ok {
				if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
					exitCode = status.ExitStatus()
				} else {
					exitCode = 1
				}
			} else {
				exitCode = -1
				stderr.WriteString("\n" + err.Error())
			}
		}
		return Response{
			Stdout:   stdout.String(),
			Stderr:   stderr.String(),
			ExitCode: exitCode,
		}
	case <-time.After(ctx_timeout):
		command.Process.Kill()
		return Response{
			Stdout:   stdout.String(),
			Stderr:   "command timed out after 120s",
			ExitCode: -1,
		}
	}
}

func main() {
	key, err := loadKey()
	if err != nil {
		fmt.Fprintf(os.Stderr, "setec-agent: %s\n", err)
		os.Exit(1)
	}

	// Read one line from stdin (base64 encrypted command)
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Buffer(make([]byte, 4*1024*1024), 4*1024*1024) // 4MB buffer
	if !scanner.Scan() {
		fmt.Fprintf(os.Stderr, "setec-agent: no input\n")
		os.Exit(1)
	}

	input := scanner.Text()

	// Decrypt command
	cmdBytes, err := decrypt(key, input)
	if err != nil {
		fmt.Fprintf(os.Stderr, "setec-agent: decrypt error: %s\n", err)
		os.Exit(1)
	}

	// Execute
	resp := executeCommand(string(cmdBytes))

	// Marshal response
	respJSON, err := json.Marshal(resp)
	if err != nil {
		fmt.Fprintf(os.Stderr, "setec-agent: json error: %s\n", err)
		os.Exit(1)
	}

	// Encrypt response
	encResp, err := encrypt(key, respJSON)
	if err != nil {
		fmt.Fprintf(os.Stderr, "setec-agent: encrypt error: %s\n", err)
		os.Exit(1)
	}

	// Output
	fmt.Println(encResp)
}