feat: 支持自定义配置频道分组规则。

super321 · Jan 9, 2025 · 181df8e · 181df8e
1 parent 15b4e01
commit 181df8e
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 41 deletions.
diff --git a/cmd/iptv/cmds/channel.go b/cmd/iptv/cmds/channel.go
@@ -43,7 +43,7 @@ func NewChannelCLI() *cobra.Command {
 			// 创建IPTV客户端
 			i, err := hwctc.NewClient(&http.Client{
 				Timeout: 10 * time.Second,
-			}, conf.HWCTC, conf.Key, conf.ServerHost, conf.Headers)
+			}, conf.HWCTC, conf.Key, conf.ServerHost, conf.Headers, conf.ChGroupRulesList)
 			if err != nil {
 				return err
 			}

diff --git a/config.yml b/config.yml
@@ -16,6 +16,26 @@ headers:
   User-Agent: "Mozilla/5.0 (X11; Linux x86_64; Fhbw2.0) AppleWebKit"
   Accept-Language: "zh-CN,en-US;q=0.8"
   X-Requested-With: "com.fiberhome.iptv"
+# 频道分组规则
+# 依照顺序识别频道分组，且仅支持正则表达式
+chGroupRules:
+  - name: 央视
+    rules:
+      - "^(CCTV|中央).+?$"
+  - name: 卫视
+    rules:
+      - "^[^(热门)].+?卫视.*?$"
+  - name: 国际
+    rules:
+      - "^(CGTN|凤凰).+?$"
+  - name: 地方
+    rules:
+      - "^(SCTV|CDTV).+?$" # 四川地方频道
+      - "^(浙江|杭州|民生|钱江|教科影视|好易购|西湖|青少体育).+?$" # 浙江地方频道
+      - "^(湖北|武汉).+?$" # 湖北地方频道
+  - name: 专区
+    rules:
+      - ".+?专区$"
 
 ###############################################
 # hw平台相关设置

diff --git a/internal/app/config/config.go b/internal/app/config/config.go
@@ -2,17 +2,28 @@ package config
 
 import (
 	"errors"
+	"iptv/internal/app/iptv"
 	"iptv/internal/app/iptv/hwctc"
 	"os"
+	"regexp"
 
+	"go.uber.org/zap"
 	"gopkg.in/yaml.v3"
 )
 
+type OptionChannelGroupRules struct {
+	Name  string   `json:"name" yaml:"name"`   // 分组名称
+	Rules []string `json:"rules" yaml:"rules"` // 分组规则
+}
+
 type Config struct {
 	Key        string            `json:"key" yaml:"key"`               // 必填，8位数字，生成Authenticator的秘钥
 	ServerHost string            `json:"serverHost" yaml:"serverHost"` // 必填，HTTP请求的IPTV服务器地址端口
 	Headers    map[string]string `json:"headers" yaml:"headers"`       // 自定义HTTP请求头
 
+	OptionChGroupRulesList []OptionChannelGroupRules `json:"chGroupRules" yaml:"chGroupRules"`
+	ChGroupRulesList       []iptv.ChannelGroupRules  `json:"-" yaml:"-"` // Validate()时进行填充
+
 	HWCTC *hwctc.Config `json:"hwctc,omitempty" yaml:"hwctc,omitempty"` // hw平台相关设置
 }
 
@@ -23,6 +34,38 @@ func (c *Config) Validate() error {
 		return errors.New("invalid IPTV-Tool config")
 	}
 
+	// L()：获取全局logger
+	logger := zap.L()
+
+	// 填充频道分组的正则表达式规则
+	c.ChGroupRulesList = make([]iptv.ChannelGroupRules, 0, len(c.OptionChGroupRulesList))
+	for _, opChGroupRules := range c.OptionChGroupRulesList {
+		if opChGroupRules.Name == "" {
+			logger.Warn("The channel group name is empty. Skip it.")
+			continue
+		} else if len(opChGroupRules.Rules) == 0 {
+			logger.Warn("The channel group rule is empty. Skip it.", zap.String("groupName", opChGroupRules.Name))
+			continue
+		}
+
+		rules := make([]*regexp.Regexp, 0, len(opChGroupRules.Rules))
+		for _, ruleStr := range opChGroupRules.Rules {
+			rule, err := regexp.Compile(ruleStr)
+			if err != nil {
+				logger.Warn("The channel group rule is incorrect. Skip it.", zap.String("groupName", opChGroupRules.Name), zap.String("rule", ruleStr), zap.Error(err))
+				continue
+			}
+
+			rules = append(rules, rule)
+		}
+		if len(rules) > 0 {
+			c.ChGroupRulesList = append(c.ChGroupRulesList, iptv.ChannelGroupRules{
+				Name:  opChGroupRules.Name,
+				Rules: rules,
+			})
+		}
+	}
+
 	return nil
 }
 

diff --git a/internal/app/iptv/channel_group.go b/internal/app/iptv/channel_group.go
@@ -4,38 +4,22 @@ import (
 	"regexp"
 )
 
-const otherGroupName = "其他"
+const otherChGroupName = "其他"
 
-// chGroupRuleMap 频道分组规则
-var chGroupRuleMap = map[string][]*regexp.Regexp{
-	"央视": {
-		regexp.MustCompile("^(CCTV|中央).+?$"),
-	},
-	"卫视": {
-		regexp.MustCompile("^[^(热门)].+?卫视.*?$"),
-	},
-	"国际": {
-		regexp.MustCompile("^(CGTN|凤凰).+?$"),
-	},
-	"地方": {
-		regexp.MustCompile("^(SCTV|CDTV).+?$"), // 四川
-		regexp.MustCompile("^(浙江|杭州|民生|钱江|教科影视|好易购|西湖|青少体育).+?$"), // 浙江
-		regexp.MustCompile("^(湖北|武汉).+?$"),                        // 湖北
-	},
-	"专区": {
-		regexp.MustCompile(".+?专区$"),
-	},
+type ChannelGroupRules struct {
+	Name  string           // 分组名称
+	Rules []*regexp.Regexp // 分组规则
 }
 
 // GetChannelGroupName 根据频道名称自动获取分组名称
-func GetChannelGroupName(channelName string) string {
+func GetChannelGroupName(chGroupRulesList []ChannelGroupRules, channelName string) string {
 	// 自动识别频道的分类
-	for groupName, groupRules := range chGroupRuleMap {
-		for _, groupRule := range groupRules {
+	for _, chGroupRules := range chGroupRulesList {
+		for _, groupRule := range chGroupRules.Rules {
 			if groupRule.MatchString(channelName) {
-				return groupName
+				return chGroupRules.Name
 			}
 		}
 	}
-	return otherGroupName
+	return otherChGroupName
 }
diff --git a/internal/app/iptv/hwctc/channel.go b/internal/app/iptv/hwctc/channel.go
@@ -137,7 +137,7 @@ func (c *Client) GetAllChannelList(ctx context.Context) ([]iptv.Channel, error)
 		}
 
 		// 自动识别频道的分类
-		groupName := iptv.GetChannelGroupName(channelName)
+		groupName := iptv.GetChannelGroupName(c.chGroupRulesList, channelName)
 
 		channels = append(channels, iptv.Channel{
 			ChannelID:       string(matches[1]),

diff --git a/internal/app/iptv/hwctc/hwctc.go b/internal/app/iptv/hwctc/hwctc.go
@@ -9,11 +9,12 @@ import (
 )
 
 type Client struct {
-	httpClient *http.Client      // HTTP客户端
-	config     *Config           // hwctc相关配置
-	key        string            // 加密Authenticator的秘钥
-	originHost string            // HTTP请求的服务器地址端口
-	headers    map[string]string // 自定义HTTP请求头
+	httpClient       *http.Client             // HTTP客户端
+	config           *Config                  // hwctc相关配置
+	key              string                   // 加密Authenticator的秘钥
+	originHost       string                   // HTTP请求的服务器地址端口
+	headers          map[string]string        // 自定义HTTP请求头
+	chGroupRulesList []iptv.ChannelGroupRules // 频道分组的规则
 
 	host string // 缓存最新重定向的服务器地址和端口
 
@@ -22,7 +23,7 @@ type Client struct {
 
 var _ iptv.Client = (*Client)(nil)
 
-func NewClient(httpClient *http.Client, config *Config, key, serverHost string, headers map[string]string) (iptv.Client, error) {
+func NewClient(httpClient *http.Client, config *Config, key, serverHost string, headers map[string]string, chGroupRulesList []iptv.ChannelGroupRules) (iptv.Client, error) {
 	// config不能为空
 	if config == nil {
 		return nil, fmt.Errorf("client config is nil")
@@ -38,13 +39,14 @@ func NewClient(httpClient *http.Client, config *Config, key, serverHost string,
 	}
 
 	i := Client{
-		httpClient: httpClient,
-		config:     config,
-		key:        key,
-		originHost: serverHost,
-		headers:    headers,
-		host:       serverHost,
-		logger:     zap.L(),
+		httpClient:       httpClient,
+		config:           config,
+		key:              key,
+		originHost:       serverHost,
+		headers:          headers,
+		chGroupRulesList: chGroupRulesList,
+		host:             serverHost,
+		logger:           zap.L(),
 	}
 	if i.httpClient == nil {
 		i.httpClient = http.DefaultClient

diff --git a/internal/app/router/router.go b/internal/app/router/router.go
@@ -91,5 +91,5 @@ func newIPTVClient(conf *config.Config) (iptv.Client, error) {
 	// 创建IPTV客户端
 	return hwctc.NewClient(&http.Client{
 		Timeout: 10 * time.Second,
-	}, conf.HWCTC, conf.Key, conf.ServerHost, conf.Headers)
+	}, conf.HWCTC, conf.Key, conf.ServerHost, conf.Headers, conf.ChGroupRulesList)
 }
diff --git a/internal/pkg/util/util.go b/internal/pkg/util/util.go
@@ -3,6 +3,7 @@ package util
 import (
 	"os"
 	"path/filepath"
+	"sort"
 )
 
 // GetCurrentAbPathByExecutable 获取当前执行程序所在的绝对路径
@@ -14,3 +15,15 @@ func GetCurrentAbPathByExecutable() (string, error) {
 	res, _ := filepath.EvalSymlinks(filepath.Dir(exePath))
 	return res, nil
 }
+
+// SortedMapKeys 对Map的Key进行排序
+func SortedMapKeys[T any](maps map[string]T) []string {
+	ret := make([]string, len(maps))
+	i := 0
+	for name := range maps {
+		ret[i] = name
+		i++
+	}
+	sort.Strings(ret)
+	return ret
+}