mrs_converter.go 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. package provider
  2. import (
  3. "encoding/binary"
  4. "errors"
  5. "fmt"
  6. "io"
  7. "os"
  8. P "github.com/metacubex/mihomo/constant/provider"
  9. "github.com/klauspost/compress/zstd"
  10. )
  11. func ConvertToMrs(buf []byte, behavior P.RuleBehavior, format P.RuleFormat, w io.Writer) (err error) {
  12. strategy := newStrategy(behavior, nil)
  13. strategy, err = rulesParse(buf, strategy, format)
  14. if err != nil {
  15. return err
  16. }
  17. if strategy.Count() == 0 {
  18. return errors.New("empty rule")
  19. }
  20. if _strategy, ok := strategy.(mrsRuleStrategy); ok {
  21. if format == P.MrsRule { // export to TextRule
  22. _strategy.DumpMrs(func(key string) bool {
  23. _, err = fmt.Fprintln(w, key)
  24. if err != nil {
  25. return false
  26. }
  27. return true
  28. })
  29. return nil
  30. }
  31. var encoder *zstd.Encoder
  32. encoder, err = zstd.NewWriter(w)
  33. if err != nil {
  34. return err
  35. }
  36. defer func() {
  37. zstdErr := encoder.Close()
  38. if err == nil {
  39. err = zstdErr
  40. }
  41. }()
  42. // header
  43. _, err = encoder.Write(MrsMagicBytes[:])
  44. if err != nil {
  45. return err
  46. }
  47. // behavior
  48. _behavior := []byte{behavior.Byte()}
  49. _, err = encoder.Write(_behavior[:])
  50. if err != nil {
  51. return err
  52. }
  53. // count
  54. count := int64(_strategy.Count())
  55. err = binary.Write(encoder, binary.BigEndian, count)
  56. if err != nil {
  57. return err
  58. }
  59. // extra (reserved for future using)
  60. var extra []byte
  61. err = binary.Write(encoder, binary.BigEndian, int64(len(extra)))
  62. if err != nil {
  63. return err
  64. }
  65. _, err = encoder.Write(extra)
  66. if err != nil {
  67. return err
  68. }
  69. return _strategy.WriteMrs(encoder)
  70. } else {
  71. return ErrInvalidFormat
  72. }
  73. }
  74. func ConvertMain(args []string) {
  75. if len(args) > 3 {
  76. behavior, err := P.ParseBehavior(args[0])
  77. if err != nil {
  78. panic(err)
  79. }
  80. format, err := P.ParseRuleFormat(args[1])
  81. if err != nil {
  82. panic(err)
  83. }
  84. source := args[2]
  85. target := args[3]
  86. sourceFile, err := os.ReadFile(source)
  87. if err != nil {
  88. panic(err)
  89. }
  90. targetFile, err := os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
  91. if err != nil {
  92. panic(err)
  93. }
  94. err = ConvertToMrs(sourceFile, behavior, format, targetFile)
  95. if err != nil {
  96. panic(err)
  97. }
  98. err = targetFile.Close()
  99. if err != nil {
  100. panic(err)
  101. }
  102. } else {
  103. panic("Usage: convert-ruleset <behavior> <format> <source file> <target file>")
  104. }
  105. }