config.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  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. Manager managerConfig `toml:"manager"`
  42. Mirrors []mirrorConfig `toml:"mirrors"`
  43. }
  44. type globalConfig struct {
  45. Name string `toml:"name"`
  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 managerConfig struct {
  52. APIBase string `toml:"api_base"`
  53. CACert string `toml:"ca_cert"`
  54. Token string `toml:"token"`
  55. }
  56. type mirrorConfig struct {
  57. Name string `toml:"name"`
  58. Provider ProviderEnum `toml:"provider"`
  59. Upstream string `toml:"upstream"`
  60. Interval int `toml:"interval"`
  61. MirrorDir string `toml:"mirror_dir"`
  62. LogDir string `toml:"log_dir"`
  63. Env map[string]string `toml:"env"`
  64. Command string `toml:"command"`
  65. UseIPv6 bool `toml:"use_ipv6"`
  66. ExcludeFile string `toml:"exclude_file"`
  67. Password string `toml:"password"`
  68. Stage1Profile string `toml:"stage1_profile"`
  69. }
  70. func loadConfig(cfgFile string) (*Config, error) {
  71. if _, err := os.Stat(cfgFile); err != nil {
  72. return nil, err
  73. }
  74. cfg := new(Config)
  75. if _, err := toml.DecodeFile(cfgFile, cfg); err != nil {
  76. logger.Error(err.Error())
  77. return nil, err
  78. }
  79. return cfg, nil
  80. }