config.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. package worker
  2. import (
  3. "errors"
  4. "os"
  5. "github.com/BurntSushi/toml"
  6. )
  7. type ProviderEnum uint8
  8. const (
  9. ProvRsync ProviderEnum = iota
  10. ProvTwoStageRsync
  11. ProvCommand
  12. )
  13. func (p ProviderEnum) MarshalText() ([]byte, error) {
  14. switch p {
  15. case ProvCommand:
  16. return []byte("command"), nil
  17. case ProvRsync:
  18. return []byte("rsync"), nil
  19. case ProvTwoStageRsync:
  20. return []byte("two-stage-rsync"), nil
  21. default:
  22. return []byte{}, errors.New("Invalid ProviderEnum value")
  23. }
  24. }
  25. func (p *ProviderEnum) UnmarshalText(text []byte) error {
  26. s := string(text)
  27. switch s {
  28. case `command`:
  29. *p = ProvCommand
  30. case `rsync`:
  31. *p = ProvRsync
  32. case `two-stage-rsync`:
  33. *p = ProvTwoStageRsync
  34. default:
  35. return errors.New("Invalid value to provierEnum")
  36. }
  37. return nil
  38. }
  39. type Config struct {
  40. Global globalConfig `toml:"global"`
  41. Mirrors []mirrorConfig `toml:"mirrors"`
  42. }
  43. type globalConfig struct {
  44. Name string `toml:"name"`
  45. Token string `toml:"token"`
  46. LogDir string `toml:"log_dir"`
  47. MirrorDir string `toml:"mirror_dir"`
  48. Concurrent int `toml:"concurrent"`
  49. Interval int `toml:"interval"`
  50. }
  51. type mirrorConfig struct {
  52. Name string `toml:"name"`
  53. Provider ProviderEnum `toml:"provider"`
  54. Upstream string `toml:"upstream"`
  55. Interval int `toml:"interval"`
  56. MirrorDir string `toml:"mirror_dir"`
  57. LogDir string `toml:"log_dir"`
  58. Env map[string]string `toml:"env"`
  59. Command string `toml:"command"`
  60. UseIPv6 bool `toml:"use_ipv6"`
  61. ExcludeFile string `toml:"exclude_file"`
  62. Password string `toml:"password"`
  63. Stage1Profile string `toml:"stage1_profile"`
  64. }
  65. func loadConfig(cfgFile string) (*Config, error) {
  66. if _, err := os.Stat(cfgFile); err != nil {
  67. return nil, err
  68. }
  69. cfg := new(Config)
  70. if _, err := toml.DecodeFile(cfgFile, cfg); err != nil {
  71. logger.Error(err.Error())
  72. return nil, err
  73. }
  74. return cfg, nil
  75. }