// security-agent is the deliberately tiny privileged boundary for manual blocking.
// It accepts only a fixed JSON protocol over a Unix socket; it never invokes a shell
// and never accepts command, jail, action, path, or argument strings from NestJS.
package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type request struct {
	Action          string `json:"action"`
	OperationKey    string `json:"operationKey"`
	SourceIP        string `json:"sourceIp"`
	Executor        string `json:"executor"`
	DurationSeconds int    `json:"durationSeconds"`
	Version         int    `json:"version"`
	Rules           []rule `json:"rules"`
}
type rule struct {
	Code            string `json:"code"`
	Enabled         bool   `json:"enabled"`
	Threshold       int    `json:"threshold"`
	WindowSeconds   int    `json:"windowSeconds"`
	CooldownSeconds int    `json:"cooldownSeconds"`
}
type response struct {
	OK        bool   `json:"ok"`
	Reference string `json:"reference,omitempty"`
	Blocked   bool   `json:"blocked,omitempty"`
	Active    bool   `json:"active,omitempty"`
	Error     string `json:"error,omitempty"`
}
type block struct {
	OperationKey string    `json:"operationKey"`
	SourceIP     string    `json:"sourceIp"`
	Executor     string    `json:"executor"`
	ExpiresAt    time.Time `json:"expiresAt"`
}
type state struct {
	Blocks      map[string]block `json:"blocks"`
	RuleVersion int              `json:"ruleVersion"`
}

var allowedDurations = map[int]bool{600: true, 3600: true, 86400: true, 604800: true}
var allowedRules = map[string]bool{"admin_login_failure": true, "client_login_failure": true, "ssh_auth_failure": true, "cmpp_auth_failure": true, "cmpp_protocol_abuse": true, "http_invalid_api_key": true, "http_signature_failure": true, "http_replay_attempt": true, "http_malicious_scan": true}

type agent struct {
	mu                                      sync.Mutex
	statePath, nginxInclude, fail2banConfig string
	data                                    state
}

func main() {
	if len(os.Args) == 4 && os.Args[1] == "report" {
		if err := reportEvent(os.Args[2], os.Args[3]); err != nil {
			log.Fatal(err)
		}
		return
	}
	socketPath := env("SECURITY_AGENT_SOCKET", "/run/cmpp-security-agent/agent.sock")
	a := &agent{statePath: env("SECURITY_AGENT_STATE", "/var/lib/cmpp-security-agent/state.json"), nginxInclude: env("SECURITY_NGINX_DENY_INCLUDE", "/etc/nginx/snippets/cmpp-security-deny.conf"), fail2banConfig: env("SECURITY_FAIL2BAN_CONFIG", "/etc/fail2ban/jail.d/cmpp-platform-generated.local"), data: state{Blocks: map[string]block{}}}
	if err := a.load(); err != nil {
		log.Fatalf("load state: %v", err)
	}
	if err := os.MkdirAll(filepath.Dir(socketPath), 0750); err != nil {
		log.Fatal(err)
	}
	_ = os.Remove(socketPath)
	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		log.Fatal(err)
	}
	if err := os.Chmod(socketPath, 0660); err != nil {
		log.Fatal(err)
	}
	defer listener.Close()
	log.Printf("security agent listening on %s", socketPath)
	for {
		connection, err := listener.Accept()
		if err != nil {
			log.Printf("accept: %v", err)
			continue
		}
		go a.serve(connection)
	}
}

func (a *agent) serve(connection net.Conn) {
	defer connection.Close()
	_ = connection.SetDeadline(time.Now().Add(5 * time.Second))
	var req request
	if err := json.NewDecoder(bufio.NewReader(connection)).Decode(&req); err != nil {
		write(connection, response{Error: "invalid request"})
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	a.prune()
	var result response
	switch req.Action {
	case "block":
		result = a.block(req)
	case "unblock":
		result = a.unblock(req)
	case "status":
		result = a.status(req)
	case "apply_rules":
		result = a.applyRules(req)
	default:
		result = response{Error: "unsupported action"}
	}
	write(connection, result)
}

func (a *agent) block(req request) response {
	if net.ParseIP(req.SourceIP) == nil || !allowedDurations[req.DurationSeconds] || (req.Executor != "nftables" && req.Executor != "nginx_real_ip") || len(req.OperationKey) < 16 {
		return response{Error: "invalid fixed block parameters"}
	}
	if existing, ok := a.data.Blocks[req.OperationKey]; ok {
		return response{OK: true, Reference: req.OperationKey, Blocked: existing.ExpiresAt.After(time.Now())}
	}
	if req.Executor == "nftables" {
		if err := nftBlock(req.SourceIP, req.DurationSeconds); err != nil {
			return response{Error: err.Error()}
		}
	}
	a.data.Blocks[req.OperationKey] = block{OperationKey: req.OperationKey, SourceIP: req.SourceIP, Executor: req.Executor, ExpiresAt: time.Now().Add(time.Duration(req.DurationSeconds) * time.Second)}
	if req.Executor == "nginx_real_ip" {
		if err := a.writeNginx(); err != nil {
			delete(a.data.Blocks, req.OperationKey)
			return response{Error: err.Error()}
		}
	}
	if err := a.save(); err != nil {
		return response{Error: err.Error()}
	}
	return a.status(req)
}

func (a *agent) unblock(req request) response {
	if net.ParseIP(req.SourceIP) == nil || (req.Executor != "nftables" && req.Executor != "nginx_real_ip") {
		return response{Error: "invalid fixed unblock parameters"}
	}
	if req.Executor == "nftables" {
		if err := nftUnblock(req.SourceIP); err != nil {
			return response{Error: err.Error()}
		}
	}
	for key, item := range a.data.Blocks {
		if item.SourceIP == req.SourceIP && item.Executor == req.Executor {
			delete(a.data.Blocks, key)
		}
	}
	if req.Executor == "nginx_real_ip" {
		if err := a.writeNginx(); err != nil {
			return response{Error: err.Error()}
		}
	}
	if err := a.save(); err != nil {
		return response{Error: err.Error()}
	}
	return response{OK: true, Reference: req.OperationKey}
}

func (a *agent) status(req request) response {
	active := command("systemctl", "is-active", "--quiet", "fail2ban") == nil
	if req.SourceIP == "" {
		return response{OK: true, Active: active}
	}
	blocked := false
	if req.Executor == "nftables" {
		family := "blocked_ipv6"
		if net.ParseIP(req.SourceIP).To4() != nil {
			family = "blocked_ipv4"
		}
		output, err := exec.Command("nft", "list", "set", "inet", "cmpp_security", family).CombinedOutput()
		if err != nil {
			return response{Error: "nftables readback failed: " + strings.TrimSpace(string(output))}
		}
		blocked = strings.Contains(string(output), req.SourceIP+" timeout") || strings.Contains(string(output), req.SourceIP+" expires")
	} else if req.Executor == "nginx_real_ip" {
		content, err := os.ReadFile(a.nginxInclude)
		if err != nil {
			return response{Error: "nginx deny readback failed: " + err.Error()}
		}
		blocked = strings.Contains(string(content), "deny "+req.SourceIP+";")
	} else {
		return response{Error: "invalid executor"}
	}
	for _, item := range a.data.Blocks {
		if blocked && item.SourceIP == req.SourceIP && item.Executor == req.Executor && item.ExpiresAt.After(time.Now()) {
			return response{OK: true, Active: active, Blocked: true, Reference: item.OperationKey}
		}
	}
	return response{OK: true, Active: active, Blocked: false}
}

func (a *agent) applyRules(req request) response {
	if req.Version <= a.data.RuleVersion {
		return response{OK: true, Reference: strconv.Itoa(a.data.RuleVersion)}
	}
	for _, item := range req.Rules {
		if !allowedRules[item.Code] || item.Threshold < 1 || item.Threshold > 100000 || item.WindowSeconds < 10 || item.WindowSeconds > 86400 || item.CooldownSeconds < 0 || item.CooldownSeconds > 604800 {
			return response{Error: "invalid fixed rule configuration"}
		}
	}
	var ssh, scan *rule
	for index := range req.Rules {
		if req.Rules[index].Code == "ssh_auth_failure" {
			ssh = &req.Rules[index]
		}
		if req.Rules[index].Code == "http_malicious_scan" {
			scan = &req.Rules[index]
		}
	}
	content := "# Generated by cmpp-security-agent. Manual changes will be overwritten.\n"
	if ssh != nil {
		content += jail("sshd", *ssh)
	}
	if scan != nil {
		content += jail("cmpp-http-scan", *scan)
	}
	previous, previousErr := os.ReadFile(a.fail2banConfig)
	if err := atomicWrite(a.fail2banConfig, []byte(content), 0640); err != nil {
		return response{Error: err.Error()}
	}
	if err := command("fail2ban-client", "-t"); err != nil {
		restore(a.fail2banConfig, previous, previousErr, 0640)
		return response{Error: "fail2ban validation failed: " + err.Error()}
	}
	if err := command("fail2ban-client", "reload"); err != nil {
		restore(a.fail2banConfig, previous, previousErr, 0640)
		return response{Error: "fail2ban reload failed: " + err.Error()}
	}
	a.data.RuleVersion = req.Version
	if err := a.save(); err != nil {
		return response{Error: err.Error()}
	}
	return response{OK: true, Reference: strconv.Itoa(req.Version), Active: true}
}

func jail(name string, item rule) string {
	enabled := "false"
	if item.Enabled {
		enabled = "true"
	}
	bantime := item.CooldownSeconds
	if bantime < item.WindowSeconds {
		bantime = item.WindowSeconds
	}
	return fmt.Sprintf("\n[%s]\nenabled = %s\nfindtime = %d\nmaxretry = %d\nbantime = %d\naction = cmpp-report-only\n", name, enabled, item.WindowSeconds, item.Threshold, bantime)
}
func nftBlock(ip string, seconds int) error {
	family := "blocked_ipv6"
	if net.ParseIP(ip).To4() != nil {
		family = "blocked_ipv4"
	}
	return command("nft", "add", "element", "inet", "cmpp_security", family, fmt.Sprintf("{ %s timeout %ds }", ip, seconds))
}
func nftUnblock(ip string) error {
	family := "blocked_ipv6"
	if net.ParseIP(ip).To4() != nil {
		family = "blocked_ipv4"
	}
	err := command("nft", "delete", "element", "inet", "cmpp_security", family, fmt.Sprintf("{ %s }", ip))
	if err != nil && strings.Contains(err.Error(), "No such file") {
		return nil
	}
	return err
}
func (a *agent) writeNginx() error {
	ips := []string{}
	for _, item := range a.data.Blocks {
		if item.Executor == "nginx_real_ip" && item.ExpiresAt.After(time.Now()) {
			ips = append(ips, item.SourceIP)
		}
	}
	sort.Strings(ips)
	lines := []string{"# Generated by cmpp-security-agent."}
	for _, ip := range ips {
		lines = append(lines, "deny "+ip+";")
	}
	previous, previousErr := os.ReadFile(a.nginxInclude)
	if err := atomicWrite(a.nginxInclude, []byte(strings.Join(lines, "\n")+"\n"), 0640); err != nil {
		return err
	}
	if err := command("nginx", "-t"); err != nil {
		restore(a.nginxInclude, previous, previousErr, 0640)
		return err
	}
	if err := command("systemctl", "reload", "nginx"); err != nil {
		restore(a.nginxInclude, previous, previousErr, 0640)
		_ = command("systemctl", "reload", "nginx")
		return err
	}
	return nil
}
func (a *agent) prune() {
	changed := false
	for key, item := range a.data.Blocks {
		if !item.ExpiresAt.After(time.Now()) {
			delete(a.data.Blocks, key)
			changed = true
		}
	}
	if changed {
		_ = a.writeNginx()
		_ = a.save()
	}
}
func (a *agent) load() error {
	bytes, err := os.ReadFile(a.statePath)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	return json.Unmarshal(bytes, &a.data)
}
func (a *agent) save() error {
	bytes, err := json.MarshalIndent(a.data, "", "  ")
	if err != nil {
		return err
	}
	return atomicWrite(a.statePath, bytes, 0600)
}
func atomicWrite(path string, bytes []byte, mode os.FileMode) error {
	if err := os.MkdirAll(filepath.Dir(path), 0750); err != nil {
		return err
	}
	temporary := path + ".tmp"
	if err := os.WriteFile(temporary, bytes, mode); err != nil {
		return err
	}
	return os.Rename(temporary, path)
}
func restore(path string, previous []byte, previousErr error, mode os.FileMode) {
	if previousErr == nil {
		_ = atomicWrite(path, previous, mode)
	} else if errors.Is(previousErr, os.ErrNotExist) {
		_ = os.Remove(path)
	}
}
func command(name string, args ...string) error {
	output, err := exec.Command(name, args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("%s: %s", err, strings.TrimSpace(string(output)))
	}
	return nil
}
func write(connection net.Conn, value response) { _ = json.NewEncoder(connection).Encode(value) }
func env(name, fallback string) string {
	if value := strings.TrimSpace(os.Getenv(name)); value != "" {
		return value
	}
	return fallback
}
func reportEvent(jail, ip string) error {
	ruleCode := map[string]string{"sshd": "ssh_auth_failure", "cmpp-http-scan": "http_malicious_scan"}[jail]
	if ruleCode == "" || net.ParseIP(ip) == nil {
		return errors.New("unsupported report event")
	}
	body := fmt.Sprintf(`{"ruleCode":%q,"sourceIp":%q,"protocol":%q,"resultCode":%q}`, ruleCode, ip, "fail2ban", jail)
	client := &http.Client{Timeout: 3 * time.Second}
	request, err := http.NewRequest(http.MethodPost, env("SECURITY_EVENT_URL", "http://127.0.0.1:3000/api/gateway/events/security-detection"), strings.NewReader(body))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("X-Security-Event-Token", os.Getenv("SECURITY_EVENT_TOKEN"))
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return fmt.Errorf("event collector returned %s", response.Status)
	}
	return nil
}