Prism/internal/subscription/service.go

package subscription

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"gopkg.in/yaml.v3"

	"github.com/prism/proxy/internal/country"
	"github.com/prism/proxy/internal/store"
)

const subscriptionUserAgent = "clash.meta/1.19.0"

type Service struct {
	store  *store.Store
	client *http.Client
}

func NewService(s *store.Store) *Service {
	return &Service{
		store: s,
		client: &http.Client{
			Timeout: 30 * time.Second,
			CheckRedirect: func(req *http.Request, via []*http.Request) error {
				if len(via) >= 10 {
					return fmt.Errorf("too many redirects")
				}
				req.Header.Set("User-Agent", subscriptionUserAgent)
				return nil
			},
		},
	}
}

type ParsedData struct {
	Nodes  []store.ProxyNode
	Groups []store.ProxyGroup
	Rules  []store.Rule
}

func (s *Service) FetchAndStore(ctx context.Context, subID int64) error {
	sub, err := s.store.GetSubscription(ctx, subID)
	if err != nil {
		return err
	}

	body, err := s.fetchURL(sub.URL)
	if err != nil {
		_ = s.store.UpdateSubscriptionFetch(ctx, subID, 0, err.Error())
		return err
	}

	parsed, err := s.parseSubscription(body, subID, sub.Type)
	if err != nil {
		_ = s.store.UpdateSubscriptionFetch(ctx, subID, 0, err.Error())
		return err
	}

	if err := s.store.ReplaceSubscriptionData(ctx, subID, parsed.Nodes, parsed.Groups, parsed.Rules); err != nil {
		_ = s.store.UpdateSubscriptionFetch(ctx, subID, 0, err.Error())
		return err
	}
	return s.store.UpdateSubscriptionFetch(ctx, subID, len(parsed.Nodes), "")
}

func (s *Service) ParseSubscription(body []byte, subID int64, subType string) (ParsedData, error) {
	return s.parseSubscription(body, subID, subType)
}

func (s *Service) parseSubscription(body []byte, subID int64, subType string) (ParsedData, error) {
	content := decodeMaybeBase64(body)
	text := strings.TrimSpace(string(content))

	format := subType
	if format == "" || format == "auto" {
		format = detectFormat(text)
	}

	switch format {
	case "links":
		return s.parseLinks(body, subID)
	default:
		parsed, err := s.parseClash(body, subID)
		if err != nil && looksLikeLinks(text) {
			return s.parseLinks(body, subID)
		}
		return parsed, err
	}
}

func detectFormat(text string) string {
	if strings.HasPrefix(text, "port:") || strings.HasPrefix(text, "mixed-port:") ||
		strings.HasPrefix(text, "proxies:") || strings.Contains(text, "\nproxies:") {
		return "clash"
	}
	if looksLikeLinks(text) {
		return "links"
	}
	return "clash"
}

func looksLikeLinks(text string) bool {
	lower := strings.ToLower(text)
	for _, prefix := range []string{"ss://", "vmess://", "trojan://", "socks5://", "http://", "https://"} {
		if strings.Contains(lower, prefix) {
			return true
		}
	}
	return false
}

func (s *Service) fetchURL(rawURL string) ([]byte, error) {
	req, err := http.NewRequest(http.MethodGet, rawURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", subscriptionUserAgent)
	req.Header.Set("Accept", "*/*")
	resp, err := s.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("fetch failed: HTTP %d", resp.StatusCode)
	}
	return io.ReadAll(resp.Body)
}

func (s *Service) parseClash(body []byte, subID int64) (ParsedData, error) {
	content := decodeMaybeBase64(body)

	var raw map[string]any
	if err := yaml.Unmarshal(content, &raw); err != nil {
		return ParsedData{}, fmt.Errorf("invalid clash yaml: %w", err)
	}
	if _, ok := raw["proxies"]; !ok {
		return ParsedData{}, fmt.Errorf("clash config missing proxies")
	}

	var out ParsedData

	if proxies, ok := raw["proxies"].([]any); ok {
		for i, p := range proxies {
			pm, ok := p.(map[string]any)
			if !ok {
				continue
			}
			name, _ := pm["name"].(string)
			if name == "" {
				name = fmt.Sprintf("node%d", i+1)
			}
			rawJSON, _ := json.Marshal(pm)
			out.Nodes = append(out.Nodes, store.ProxyNode{
				Name:        name,
				RuntimeName: store.RuntimeName("sub", subID, fmt.Sprintf("%d_%s", i, name)),
				Type:        fmt.Sprint(pm["type"]),
				RawConfig:   string(rawJSON),
				Country:     country.Detect(name),
				Enabled:     true,
			})
		}
	}

	if groups, ok := raw["proxy-groups"].([]any); ok {
		for i, g := range groups {
			gm, ok := g.(map[string]any)
			if !ok {
				continue
			}
			name, _ := gm["name"].(string)
			if name == "" {
				continue
			}
			gtype, _ := gm["type"].(string)
			var members []string
			if arr, ok := gm["proxies"].([]any); ok {
				for _, m := range arr {
					members = append(members, fmt.Sprint(m))
				}
			}
			membersJSON, _ := json.Marshal(members)
			out.Groups = append(out.Groups, store.ProxyGroup{
				Name:        name,
				RuntimeName: store.RuntimeName("grp", subID, fmt.Sprintf("%d_%s", i, name)),
				Type:        gtype,
				Proxies:     string(membersJSON),
				Enabled:     true,
			})
		}
	}

	if rules, ok := raw["rules"].([]any); ok {
		for _, r := range rules {
			line := fmt.Sprint(r)
			parts := strings.Split(line, ",")
			if len(parts) < 2 {
				continue
			}
			rule := store.Rule{
				RuleType: parts[0],
				Target:   parts[len(parts)-1],
				Enabled:  true,
			}
			if len(parts) > 2 {
				rule.Payload = strings.Join(parts[1:len(parts)-1], ",")
			}
			out.Rules = append(out.Rules, rule)
		}
	}

	return out, nil
}

func (s *Service) parseLinks(body []byte, subID int64) (ParsedData, error) {
	text := string(decodeMaybeBase64(body))
	lines := strings.Split(text, "\n")
	var out ParsedData
	idx := 0
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		node, err := parseLink(line, subID, idx)
		if err != nil {
			continue
		}
		out.Nodes = append(out.Nodes, node)
		idx++
	}
	if len(out.Nodes) == 0 {
		return ParsedData{}, fmt.Errorf("no valid proxy links found")
	}
	names := make([]string, 0, len(out.Nodes))
	for _, n := range out.Nodes {
		names = append(names, n.RuntimeName)
	}
	namesJSON, _ := json.Marshal(names)
	out.Groups = append(out.Groups, store.ProxyGroup{
		Name:        fmt.Sprintf("sub%d", subID),
		RuntimeName: store.RuntimeName("grp", subID, "auto"),
		Type:        "url-test",
		Proxies:     string(namesJSON),
		Enabled:     true,
	})
	return out, nil
}

func parseLink(link string, subID int64, idx int) (store.ProxyNode, error) {
	lower := strings.ToLower(link)
	var proxy map[string]any
	var name string
	var ptype string

	switch {
	case strings.HasPrefix(lower, "ss://"):
		proxy, name, ptype = parseSS(link)
	case strings.HasPrefix(lower, "socks5://"):
		proxy, name, ptype = parseSocks5(link)
	case strings.HasPrefix(lower, "http://"), strings.HasPrefix(lower, "https://"):
		proxy, name, ptype = parseHTTP(link)
	case strings.HasPrefix(lower, "trojan://"):
		proxy, name, ptype = parseTrojan(link)
	default:
		return store.ProxyNode{}, fmt.Errorf("unsupported link")
	}
	if proxy == nil {
		return store.ProxyNode{}, fmt.Errorf("failed to parse link")
	}
	if name == "" {
		name = fmt.Sprintf("node%d", idx+1)
	}
	raw, _ := json.Marshal(proxy)
	return store.ProxyNode{
		Name:        name,
		RuntimeName: store.RuntimeName("sub", subID, name),
		Type:        ptype,
		RawConfig:   string(raw),
		Country:     country.Detect(name),
		Enabled:     true,
	}, nil
}

func parseSS(link string) (map[string]any, string, string) {
	name, payload := splitLinkFragment(link, "ss://")
	if payload == "" {
		return nil, name, "ss"
	}

	// SIP002: ss://BASE64(method:password@host:port)#name
	if proxy, ok := parseSIP002(payload); ok {
		return proxy, name, "ss"
	}

	// 传统格式: ss://BASE64(method:password)@host:port#name
	// 或 ss://BASE64(method:password)@BASE64(host:port)#name（部分机场订阅）
	parts := strings.SplitN(payload, "@", 2)
	if len(parts) != 2 {
		return nil, name, "ss"
	}
	userinfo := string(decodeBase64Flexible(parts[0]))
	if userinfo == "" {
		userinfo = parts[0]
	}
	hostport := string(decodeBase64Flexible(parts[1]))
	if hostport == "" {
		hostport = parts[1]
	}
	server, portStr, ok := strings.Cut(hostport, ":")
	if !ok {
		return nil, name, "ss"
	}
	var port int
	fmt.Sscanf(portStr, "%d", &port)
	cipher, password, _ := strings.Cut(userinfo, ":")
	if cipher == "" {
		cipher = "aes-256-gcm"
	}
	return map[string]any{
		"type":     "ss",
		"server":   server,
		"port":     port,
		"cipher":   cipher,
		"password": password,
	}, name, "ss"
}

func parseSIP002(payload string) (map[string]any, bool) {
	decoded := string(decodeBase64Flexible(payload))
	at := strings.LastIndex(decoded, "@")
	if at <= 0 || !strings.Contains(decoded[:at], ":") {
		return nil, false
	}
	userinfo := decoded[:at]
	hostport := decoded[at+1:]
	server, portStr, ok := strings.Cut(hostport, ":")
	if !ok || server == "" {
		return nil, false
	}
	var port int
	if _, err := fmt.Sscanf(portStr, "%d", &port); err != nil || port <= 0 {
		return nil, false
	}
	cipher, password, ok := strings.Cut(userinfo, ":")
	if !ok || cipher == "" {
		return nil, false
	}
	return map[string]any{
		"type":     "ss",
		"server":   server,
		"port":     port,
		"cipher":   cipher,
		"password": password,
	}, true
}

func splitLinkFragment(link, scheme string) (name, payload string) {
	rest := strings.TrimPrefix(link, scheme)
	if i := strings.Index(rest, "#"); i >= 0 {
		name, _ = url.QueryUnescape(rest[i+1:])
		rest = rest[:i]
	}
	return name, rest
}

func decodeBase64Flexible(s string) []byte {
	decoders := []func(string) ([]byte, error){
		base64.StdEncoding.DecodeString,
		base64.RawStdEncoding.DecodeString,
		base64.URLEncoding.DecodeString,
		base64.RawURLEncoding.DecodeString,
	}
	for _, dec := range decoders {
		if b, err := dec(s); err == nil && len(b) > 0 {
			return b
		}
	}
	return nil
}

func parseSocks5(link string) (map[string]any, string, string) {
	name, payload := splitLinkFragment(link, "socks5://")
	user, pass := "", ""
	if at := strings.LastIndex(payload, "@"); at >= 0 {
		auth := payload[:at]
		payload = payload[at+1:]
		if c := strings.Index(auth, ":"); c >= 0 {
			user, pass = auth[:c], auth[c+1:]
		}
	}
	server, portStr, _ := strings.Cut(payload, ":")
	var port int
	fmt.Sscanf(portStr, "%d", &port)
	m := map[string]any{"type": "socks5", "server": server, "port": port}
	if user != "" {
		m["username"] = user
		m["password"] = pass
	}
	return m, name, "socks5"
}

func parseHTTP(link string) (map[string]any, string, string) {
	isTLS := strings.HasPrefix(strings.ToLower(link), "https://")
	scheme := "http://"
	if isTLS {
		scheme = "https://"
	}
	name, payload := splitLinkFragment(link, scheme)
	user, pass := "", ""
	if at := strings.LastIndex(payload, "@"); at >= 0 {
		auth := payload[:at]
		payload = payload[at+1:]
		if c := strings.Index(auth, ":"); c >= 0 {
			user, pass = auth[:c], auth[c+1:]
		}
	}
	server, portStr, _ := strings.Cut(payload, ":")
	var port int
	fmt.Sscanf(portStr, "%d", &port)
	m := map[string]any{"type": "http", "server": server, "port": port, "tls": isTLS}
	if user != "" {
		m["username"] = user
		m["password"] = pass
	}
	return m, name, "http"
}

func parseTrojan(link string) (map[string]any, string, string) {
	name, payload := splitLinkFragment(link, "trojan://")
	pass, rest, _ := strings.Cut(payload, "@")
	server, portStr, _ := strings.Cut(rest, ":")
	var port int
	fmt.Sscanf(portStr, "%d", &port)
	if q := strings.Index(server, "?"); q >= 0 {
		server = server[:q]
	}
	return map[string]any{
		"type":     "trojan",
		"server":   server,
		"port":     port,
		"password": pass,
		"sni":      server,
	}, name, "trojan"
}

func decodeMaybeBase64(body []byte) []byte {
	trim := strings.TrimSpace(string(body))
	if strings.HasPrefix(trim, "proxies:") || strings.HasPrefix(trim, "{") || strings.HasPrefix(trim, "ss://") {
		return body
	}
	if decoded := decodeBase64Flexible(trim); len(decoded) > 0 {
		return decoded
	}
	return body
}