-
Notifications
You must be signed in to change notification settings - Fork 1
/
root.go
118 lines (98 loc) · 3.22 KB
/
root.go
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 cmd
import (
"fmt"
"os"
"os/signal"
"syscall"
"time"
"github.com/axllent/bwlog/app"
"github.com/axllent/bwlog/utils"
"github.com/spf13/cobra"
)
// rootCmd represents the base command when called without any subcommands
var rootCmd = &cobra.Command{
Use: "bwlog -i eth0 -d ~/bwlog/",
Short: "BWLog: A lightweight bandwidth logger",
SilenceErrors: true, // suppress duplicate error on error
Args: cobra.ExactArgs(0),
RunE: func(cmd *cobra.Command, args []string) error {
if err := initConfig(cmd); err != nil {
return err
}
for _, nwIf := range app.Config.Interfaces {
app.LoadStats(nwIf)
}
app.SyncNwInterfaces()
sigs := make(chan os.Signal, 1)
// catch all signals since not explicitly listing
// Program that will listen to the SIGINT and SIGTERM
// SIGINT will listen to CTRL-C.
// SIGTERM will be caught if kill command executed
signal.Notify(sigs, os.Interrupt, syscall.SIGTERM)
// method invoked upon seeing signal
go func() {
s := <-sigs
fmt.Printf("Got %s signal, saving data & shutting down...\n", s)
app.SyncNwInterfaces()
app.SaveStats()
os.Exit(0)
}()
app.StartHTTP()
ticker := time.NewTicker(app.Config.SaveInterval)
for range ticker.C {
app.SyncNwInterfaces()
app.SaveStats()
}
return nil
},
}
// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
if err := rootCmd.Execute(); err != nil {
fmt.Println(err)
os.Exit(1)
}
}
func init() {
rootCmd.Flags().StringVarP(&app.Config.DatabaseDir, "database", "d", "./", "Database directory to save CSV files")
rootCmd.Flags().StringVarP(&app.Config.Listen, "listen", "l", "0.0.0.0:8080", "Interface & port to listen on")
rootCmd.Flags().StringVar(&app.Config.SSLCert, "sslcert", "", "SSL certificate (must be used together with --sslkey)")
rootCmd.Flags().StringVar(&app.Config.SSLKey, "sslkey", "", "SSL key (must be used together with --sslcert)")
rootCmd.Flags().StringP("interfaces", "i", "", "Interfaces to monitor, comma separated eg: eth0,eth1")
rootCmd.Flags().StringP("password", "p", "", "Auth password file (must contain a single \"<user> <pass>\")")
rootCmd.Flags().StringP("save", "s", "60s", "How often to save the database to disk. Examples: 30s, 5m, 1h")
rootCmd.MarkFlagRequired("interfaces")
// hide `-h`
rootCmd.Flags().BoolP("help", "h", false, "override help so we can hide it")
rootCmd.Flags().MarkHidden("help")
// hide `help` subcommand
rootCmd.SetHelpCommand(&cobra.Command{
Use: "no-help",
Hidden: true,
})
}
// initConfig reads in config file and ENV variables if set.
func initConfig(cmd *cobra.Command) error {
if err := app.InitInterfaces(cmd); err != nil {
return err
}
passwordFile, _ := cmd.Flags().GetString("password")
if passwordFile != "" {
if err := app.BasicAuthFromFile(passwordFile); err != nil {
return err
}
}
save, _ := cmd.Flags().GetString("save")
intrvl, err := time.ParseDuration(save)
if err != nil {
return err
}
app.Config.SaveInterval = intrvl
if !utils.IsDir(app.Config.DatabaseDir) {
if err := os.Mkdir(app.Config.DatabaseDir, 0755); err != nil {
return err
}
}
return nil
}