Prism/internal/configbuilder/builder.go

package configbuilder

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"

	"github.com/prism/proxy/internal/country"
	"github.com/prism/proxy/internal/geodata"
	"github.com/prism/proxy/internal/store"
	"gopkg.in/yaml.v3"
)

type Builder struct {
	store   *store.Store
	dataDir string
}

func New(s *store.Store, dataDir string) *Builder {
	return &Builder{store: s, dataDir: dataDir}
}

func (b *Builder) Build(ctx context.Context) ([]byte, error) {
	return b.build(ctx, false)
}

// BuildMinimal loads nodes/groups with MATCH-only rules when full config fails.
func (b *Builder) BuildMinimal(ctx context.Context) ([]byte, error) {
	return b.build(ctx, true)
}

func (b *Builder) build(ctx context.Context, minimalRules bool) ([]byte, error) {
	settings, err := b.store.GetSettings(ctx)
	if err != nil {
		return nil, err
	}

	httpPort := settings["http_port"]
	socksPort := settings["socks_port"]
	logLevel := settings["log_level"]
	if logLevel == "" {
		logLevel = "info"
	}
	defaultPolicy := settings["default_policy"]
	if defaultPolicy == "" {
		defaultPolicy = "DIRECT"
	}

	nodes, err := b.store.ListProxyNodes(ctx, "", nil)
	if err != nil {
		return nil, err
	}
	groups, err := b.store.ListProxyGroups(ctx)
	if err != nil {
		return nil, err
	}
	mergedRules, err := b.store.ListMergedRules(ctx, defaultPolicy)
	if err != nil {
		return nil, err
	}
	if minimalRules {
		mergedRules = nil
	}

	cfg := map[string]any{
		"port":       atoi(httpPort, 7890),
		"socks-port": atoi(socksPort, 7891),
		"allow-lan":  true,
		"mode":       "rule",
		"log-level":  logLevel,
		"proxies":    []any{},
		"proxy-groups": []any{
			map[string]any{
				"name":    "PROXY",
				"type":    "select",
				"proxies": []string{"DIRECT"},
			},
		},
		"rules": []string{"MATCH," + defaultPolicy},
	}
	applyInboundAuth(cfg, settings)

	var proxies []any
	proxyNames := []string{}
	for _, n := range nodes {
		if !n.Enabled {
			continue
		}
		var proxy map[string]any
		if err := json.Unmarshal([]byte(n.RawConfig), &proxy); err != nil {
			continue
		}
		proxy["name"] = n.RuntimeName
		proxies = append(proxies, proxy)
		proxyNames = append(proxyNames, n.RuntimeName)
	}
	cfg["proxies"] = proxies

	var proxyGroups []any
	hasGroups := false
	for _, g := range groups {
		if !g.Enabled {
			continue
		}
		hasGroups = true
		var members []string
		_ = json.Unmarshal([]byte(g.Proxies), &members)
		remapped := remapMembers(members, nodes, groups)
		proxyGroups = append(proxyGroups, buildProxyGroup(g.RuntimeName, g.Type, remapped))
	}
	if !hasGroups && len(proxyNames) > 0 {
		proxyGroups = []any{
			map[string]any{
				"name":    "PROXY",
				"type":    "select",
				"proxies": append(proxyNames, "DIRECT"),
			},
		}
	} else if len(proxyGroups) > 0 {
		cfg["proxy-groups"] = proxyGroups
	} else {
		cfg["proxy-groups"] = []any{
			map[string]any{
				"name":    "PROXY",
				"type":    "select",
				"proxies": []string{"DIRECT"},
			},
		}
	}
	if len(proxyGroups) > 0 {
		cfg["proxy-groups"] = proxyGroups
	}
	proxyGroups = appendCountryProxyGroups(proxyGroups, nodes)
	cfg["proxy-groups"] = proxyGroups

	nameMap := buildNameMap(nodes, groups)
	countryNameMap, countryValid := countryMaps(nodes)
	for k, v := range countryNameMap {
		nameMap[k] = v
	}
	valid := buildValidTargets(nodes, groups)
	for k := range countryValid {
		valid[k] = true
	}
	defaultPolicy = resolvePolicy(defaultPolicy, nameMap, valid, groups)
	cfg["rules"] = sanitizeRules(mergedRules, nameMap, valid, defaultPolicy, geodata.Ready(b.dataDir))

	out, err := yaml.Marshal(cfg)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func buildNameMap(nodes []store.ProxyNode, groups []store.ProxyGroup) map[string]string {
	nameMap := make(map[string]string)
	for _, n := range nodes {
		nameMap[n.Name] = n.RuntimeName
		nameMap[n.RuntimeName] = n.RuntimeName
	}
	for _, g := range groups {
		nameMap[g.Name] = g.RuntimeName
		nameMap[g.RuntimeName] = g.RuntimeName
	}
	return nameMap
}

func remapMembers(members []string, nodes []store.ProxyNode, groups []store.ProxyGroup) []string {
	nameMap := buildNameMap(nodes, groups)
	out := make([]string, 0, len(members))
	for _, name := range members {
		if name == "DIRECT" || name == "REJECT" {
			out = append(out, name)
			continue
		}
		if rn, ok := nameMap[name]; ok {
			out = append(out, rn)
		}
	}
	return out
}

func remapRuleTarget(rule string, nameMap map[string]string) string {
	parts := strings.Split(rule, ",")
	if len(parts) < 2 {
		return rule
	}
	target := parts[len(parts)-1]
	if code, ok := country.ParseTargetAlias(target); ok {
		parts[len(parts)-1] = country.GroupRuntime(code)
		return strings.Join(parts, ",")
	}
	if mapped, ok := nameMap[target]; ok {
		parts[len(parts)-1] = mapped
		return strings.Join(parts, ",")
	}
	return rule
}

func remapProxyNames(names []string, nodes []store.ProxyNode) []string {
	return remapMembers(names, nodes, nil)
}

func buildProxyGroup(name, gtype string, members []string) map[string]any {
	group := map[string]any{
		"name":    name,
		"type":    gtype,
		"proxies": appendGroupMembers(gtype, members),
	}
	switch gtype {
	case "url-test", "fallback":
		group["url"] = "http://www.gstatic.com/generate_204"
		group["interval"] = 300
	}
	return group
}

func appendGroupMembers(gtype string, members []string) []string {
	switch strings.ToLower(gtype) {
	case "url-test", "fallback":
		return members
	default:
		return append(append([]string{}, members...), "DIRECT")
	}
}

func buildValidTargets(nodes []store.ProxyNode, groups []store.ProxyGroup) map[string]bool {
	valid := map[string]bool{
		"DIRECT": true,
		"REJECT": true,
	}
	for _, n := range nodes {
		if n.Enabled {
			valid[n.RuntimeName] = true
		}
	}
	for _, g := range groups {
		if g.Enabled {
			valid[g.RuntimeName] = true
		}
	}
	return valid
}

func resolvePolicy(policy string, nameMap map[string]string, valid map[string]bool, groups []store.ProxyGroup) string {
	if code, ok := country.ParseTargetAlias(policy); ok {
		rn := country.GroupRuntime(code)
		if valid[rn] {
			return rn
		}
	}
	if mapped, ok := nameMap[policy]; ok && valid[mapped] {
		return mapped
	}
	if valid[policy] {
		return policy
	}
	if policy == "PROXY" || policy == "GLOBAL" {
		for _, g := range groups {
			if g.Enabled && valid[g.RuntimeName] {
				return g.RuntimeName
			}
		}
	}
	return "DIRECT"
}

func sanitizeRules(merged []store.MergedRule, nameMap map[string]string, valid map[string]bool, defaultPolicy string, geoReady bool) []string {
	var rules []string
	hasMatch := false
	for _, r := range merged {
		line := remapRuleTarget(r.Rule, nameMap)
		if !geoReady && needsGeoData(line) {
			continue
		}
		if strings.HasPrefix(line, "MATCH,") {
			line = "MATCH," + defaultPolicy
			hasMatch = true
		}
		target := ruleTarget(line)
		if target == "" || !valid[target] {
			continue
		}
		rules = append(rules, line)
	}
	if !hasMatch {
		rules = append(rules, "MATCH,"+defaultPolicy)
	}
	return rules
}

func needsGeoData(rule string) bool {
	upper := strings.ToUpper(rule)
	if strings.HasPrefix(upper, "GEOIP,") || strings.HasPrefix(upper, "GEOSITE,") {
		return true
	}
	return strings.Contains(strings.ToLower(rule), "geoip:")
}

func ruleTarget(line string) string {
	parts := strings.Split(line, ",")
	if len(parts) < 2 {
		return ""
	}
	return parts[len(parts)-1]
}

func applyInboundAuth(cfg map[string]any, settings map[string]string) {
	user := strings.TrimSpace(settings["proxy_auth_user"])
	pass := settings["proxy_auth_pass"]
	if user == "" || pass == "" {
		return
	}
	cfg["authentication"] = []string{user + ":" + pass}
	cfg["skip-auth-prefixes"] = []string{"127.0.0.1/8", "::1/128"}
}

func atoi(s string, def int) int {
	var v int
	if _, err := fmt.Sscanf(strings.TrimSpace(s), "%d", &v); err != nil || v <= 0 {
		return def
	}
	return v
}

func appendCountryProxyGroups(groups []any, nodes []store.ProxyNode) []any {
	byCountry := map[string][]string{}
	for _, n := range nodes {
		if !n.Enabled || n.Country == "" {
			continue
		}
		byCountry[n.Country] = append(byCountry[n.Country], n.RuntimeName)
	}
	codes := make([]string, 0, len(byCountry))
	for code := range byCountry {
		codes = append(codes, code)
	}
	sort.Strings(codes)
	for _, code := range codes {
		members := byCountry[code]
		if len(members) == 0 {
			continue
		}
		groups = append(groups, buildProxyGroup(country.GroupRuntime(code), "url-test", members))
	}
	return groups
}

func countryMaps(nodes []store.ProxyNode) (map[string]string, map[string]bool) {
	nameMap := map[string]string{}
	valid := map[string]bool{}
	seen := map[string]bool{}
	for _, n := range nodes {
		if !n.Enabled || n.Country == "" || seen[n.Country] {
			continue
		}
		seen[n.Country] = true
		rn := country.GroupRuntime(n.Country)
		nameMap[country.TargetAlias(n.Country)] = rn
		nameMap[country.Name(n.Country)] = rn
		nameMap[rn] = rn
		valid[rn] = true
	}
	return nameMap, valid
}