rss-exporter/main.go at main · 4O4-Not-F0und/rss-exporter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/sirupsen/logrus"
	"gopkg.in/yaml.v3"
)

const (
	defaultProbePath     = "/probe"
	defaultListenAddress = "0.0.0.0"
	defaultListenPort    = 9191
	defaultConfigFile    = "config.yml"
)

type Config struct {
	ListenAddress  string `yaml:"listen_address"`
	ListenPort     int    `yaml:"listen_port"`
	LogLevel       string `yaml:"log_level"`
	ProbePath      string `yaml:"probe_path"`
	DefaultTimeout int    `yaml:"default_timeout"`
}

var (
	appConfig Config
	logLevels = map[string]logrus.Level{
		"trace": logrus.TraceLevel,
		"debug": logrus.DebugLevel,
		"info":  logrus.InfoLevel,
		"warn":  logrus.WarnLevel,
	}
)

func init() {
	var configFile string
	flag.StringVar(&configFile, "config", defaultConfigFile, "path to config file")
	flag.Parse()

	logrus.SetLevel(logrus.InfoLevel)
	logrus.SetOutput(os.Stdout)
	logrus.SetFormatter(&logrus.TextFormatter{
		TimestampFormat:        time.RFC3339,
		DisableColors:          true,
		DisableLevelTruncation: true,
		ForceQuote:             true,
		FullTimestamp:          true,
	})

	var err error
	appConfig, err = loadConfig(configFile)
	if err != nil {
		logrus.Fatalf("load config failed: %v", err)
	}

	logLevel, ok := logLevels[appConfig.LogLevel]
	if ok {
		logrus.SetLevel(logLevel)
	} else {
		logrus.Warnf("Invalid log level '%s'", appConfig.LogLevel)
	}
}

func loadConfig(configFile string) (cfg Config, err error) {

	cfg.ListenAddress = defaultListenAddress
	cfg.ListenPort = defaultListenPort
	cfg.ProbePath = defaultProbePath
	cfg.DefaultTimeout = int(defaultTimeout.Seconds())

	yamlFile, err := os.ReadFile(configFile)
	if err != nil {
		if os.IsNotExist(err) {
			err = fmt.Errorf("Config file '%s' not found, starting with default config", configFile)
			return
		}
		return cfg, fmt.Errorf("load config file '%s' failed: %v", configFile, err)
	}

	err = yaml.Unmarshal(yamlFile, &cfg)
	if err != nil {
		return cfg, fmt.Errorf("parse '%s' failed: %v", configFile, err)
	}

	if cfg.ProbePath == "" {
		cfg.ProbePath = defaultProbePath
	} else if cfg.ProbePath[0] != '/' {
		cfg.ProbePath = "/" + cfg.ProbePath
	}

	logrus.Infof("loaded config from '%s'", configFile)
	return cfg, nil
}

func main() {
	http.HandleFunc(appConfig.ProbePath, probeHandler)
	http.HandleFunc("/", landingPageHandler)

	listenOn := fmt.Sprintf("%s:%d", appConfig.ListenAddress, appConfig.ListenPort)
	logrus.Infof("Listening at: %s", listenOn)
	logrus.Fatalln(http.ListenAndServe(listenOn, nil))
}

func landingPageHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	landingPageHTML := fmt.Sprintf(`<html>
			<head><title>RSS Exporter</title></head>
			<body>
			<h1>RSS Exporter</h1>
			<p>Probe Example: <code>%s?target=https://example.com/files/sample-rss-2.xml&timeout=10</code></p>
			</body>
			</html>`, appConfig.ProbePath)
	w.Write([]byte(landingPageHTML))
}