root/golib
1
0
Fork 0
mirror of https://github.com/therootcompany/golib.git synced 2026-11-03 22:44:58 +00:00
Code Issues Projects Releases Wiki Activity

feat(geoip-update): align CLI flags with conventions

Browse Source
This commit is contained in:
AJ ONeal 2026-09-08 15:19:38 -06:00
parent e82e3979d2
commit f7ca1d22a5
No known key found for this signature in database
2 changed files with 194 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

175
net/geoip/cmd/geoip-update/main.go
View File

@ -5,13 +5,19 @@ package main
import (
"context"
"encoding/base64"
"encoding/csv"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"net/http"
"os"
"os/signal"
"path/filepath"
"runtime/debug"
"strings"
"syscall"
"time"
"github.com/therootcompany/golib/net/geoip"
@ -30,7 +36,27 @@ var (
)
func printVersion(w io.Writer) {
_, _ = fmt.Fprintf(w, "%s v%s %s (%s)\n", name, version, commit[:7], date)
v, c, d := version, commit, date
if c == "0000000" {
if bi, ok := debug.ReadBuildInfo(); ok {
for _, setting := range bi.Settings {
switch setting.Key {
case "vcs.revision":
c = setting.Value
case "vcs.time":
d = setting.Value
case "vcs.modified":
if setting.Value == "true" {
d += "+dirty"
}
}
}
}
}
if len(c) > 7 {
c = c[:7]
}
_, _ = fmt.Fprintf(w, "%s v%s %s (%s)\n", name, v, c, d)
_, _ = fmt.Fprintf(w, "Copyright (C) %s %s\n", licenseYear, licenseOwner)
_, _ = fmt.Fprintf(w, "Licensed under %s\n", licenseType)
}
@ -38,10 +64,20 @@ func printVersion(w io.Writer) {
const defaultFailureBackoff = 6 * time.Hour
type Config struct {
ConfPath string
Dir string
BaseURL string
FreshDays int
ConfPath string
Dir string
BaseURL string
FreshDays int
Format string
HumanReadable bool
}
type result struct {
State string `json:"state"`
Edition string `json:"edition"`
Path string `json:"path"`
Date string `json:"date,omitempty"`
Error string `json:"error,omitempty"`
}
func main() {
@ -51,6 +87,15 @@ func main() {
fs.StringVar(&cfg.Dir, "dir", "", "directory to store .tar.gz files (overrides DatabaseDirectory in config)")
fs.StringVar(&cfg.BaseURL, "base-url", "", "base URL for database downloads (overrides DownloadBase in geoip package)")
fs.IntVar(&cfg.FreshDays, "fresh-days", 3, "skip download if file is younger than N days")
fs.StringVar(&cfg.Format, "format", "", "output format: pretty, tsv, csv, json (default: auto)")
fs.BoolVar(&cfg.HumanReadable, "human-readable", false, "use human-readable output (same as --format pretty)")
fs.BoolVar(&cfg.HumanReadable, "h", false, "use human-readable output (same as --format pretty)")
fs.Usage = func() {
w := fs.Output()
_, _ = fmt.Fprintf(w, "USAGE\n %s [flags]\n\nFLAGS\n", name)
fs.PrintDefaults()
_, _ = fmt.Fprintln(w, "\nCONFIGURATION\n GeoIP.conf uses AccountID, LicenseKey, EditionIDs, and DatabaseDirectory.")
}
if len(os.Args) > 1 {
switch os.Args[1] {
@ -69,17 +114,37 @@ func main() {
if errors.Is(err, flag.ErrHelp) {
os.Exit(0)
}
os.Exit(1)
os.Exit(2)
}
if fs.NArg() != 0 {
_, _ = fmt.Fprintf(os.Stderr, "error: unexpected arguments: %s\n", strings.Join(fs.Args(), " "))
os.Exit(2)
}
if cfg.FreshDays < 0 {
_, _ = fmt.Fprintln(os.Stderr, "error: -fresh-days must not be negative")
os.Exit(2)
}
if cfg.HumanReadable {
if cfg.Format != "" && cfg.Format != "pretty" {
_, _ = fmt.Fprintln(os.Stderr, "error: -h conflicts with a non-pretty -format")
os.Exit(2)
}
cfg.Format = "pretty"
}
format, err := resolveFormat(cfg.Format, os.Stdout)
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(2)
}
data, err := os.ReadFile(cfg.ConfPath)
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
_, _ = fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(1)
}
conf, err := geoip.ParseConf(string(data))
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
_, _ = fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(1)
}
@ -91,19 +156,22 @@ func main() {
outDir = "."
}
if err := os.MkdirAll(outDir, 0o755); err != nil {
fmt.Fprintf(os.Stderr, "error: mkdir %s: %v\n", outDir, err)
_, _ = fmt.Fprintf(os.Stderr, "error: mkdir %s: %v\n", outDir, err)
os.Exit(1)
}
if len(conf.EditionIDs) == 0 {
fmt.Fprintf(os.Stderr, "error: no EditionIDs found in %s\n", cfg.ConfPath)
_, _ = fmt.Fprintf(os.Stderr, "error: no EditionIDs found in %s\n", cfg.ConfPath)
os.Exit(1)
}
creds := base64.StdEncoding.EncodeToString([]byte(conf.AccountID + ":" + conf.LicenseKey))
authHeader := http.Header{"Authorization": []string{"Basic " + creds}}
maxAge := time.Duration(cfg.FreshDays) * 24 * time.Hour
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
results := make([]result, 0, len(conf.EditionIDs))
exitCode := 0
for _, edition := range conf.EditionIDs {
path := filepath.Join(outDir, geoip.TarGzName(edition))
@ -118,12 +186,13 @@ func main() {
cacher.MaxAge = maxAge
cacher.FailureBackoff = defaultFailureBackoff
cacher.Header = authHeader
fmt.Fprintf(os.Stderr, "Fetching %s... ", edition)
_, _ = fmt.Fprintf(os.Stderr, "Fetching %s... ", edition)
t := time.Now()
updated, err := cacher.Fetch(context.Background())
updated, err := cacher.Fetch(ctx)
if err != nil {
fmt.Fprintln(os.Stderr)
fmt.Fprintf(os.Stderr, "error: %s: %v\n", edition, err)
_, _ = fmt.Fprintln(os.Stderr)
_, _ = fmt.Fprintf(os.Stderr, "error: %s: %v\n", edition, err)
results = append(results, result{State: "error", Edition: edition, Path: path, Error: err.Error()})
exitCode = 1
continue
}
@ -131,9 +200,81 @@ func main() {
if updated {
state = "updated"
}
fmt.Fprintf(os.Stderr, "%s (%s)\n", time.Since(t).Round(time.Millisecond), state)
info, _ := os.Stat(path)
fmt.Printf("%-10s %s -> %s (%s)\n", state+":", edition, path, info.ModTime().Format("2006-01-02"))
_, _ = fmt.Fprintf(os.Stderr, "%s (%s)\n", time.Since(t).Round(time.Millisecond), state)
info, statErr := os.Stat(path)
date := ""
if statErr == nil {
date = info.ModTime().Format("2006-01-02")
}
results = append(results, result{State: state, Edition: edition, Path: path, Date: date})
}
if err := writeResults(os.Stdout, format, results); err != nil {
_, _ = fmt.Fprintf(os.Stderr, "error: write output: %v\n", err)
exitCode = 1
}
os.Exit(exitCode)
}
func resolveFormat(format string, out *os.File) (string, error) {
if format != "" {
return parseFormat(format)
}
if isTTYish(out) {
return "pretty", nil
}
return "tsv", nil
}
func parseFormat(format string) (string, error) {
switch format {
case "pretty", "tsv", "csv", "json":
return format, nil
default:
return "", fmt.Errorf("invalid format %q: want pretty, tsv, csv, or json", format)
}
}
func isTTYish(f *os.File) bool {
fi, err := f.Stat()
if err != nil {
return false
}
m := os.ModeDevice | os.ModeCharDevice
return fi.Mode()&m == m
}
func writeResults(w io.Writer, format string, results []result) error {
switch format {
case "json":
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
return enc.Encode(results)
case "csv", "tsv":
cw := csv.NewWriter(w)
if format == "tsv" {
cw.Comma = '\t'
} else if err := cw.Write([]string{"state", "edition", "path", "date", "error"}); err != nil {
return err
}
for _, r := range results {
if err := cw.Write([]string{r.State, r.Edition, r.Path, r.Date, r.Error}); err != nil {
return err
}
}
cw.Flush()
return cw.Error()
default:
for _, r := range results {
if r.State == "error" {
if _, err := fmt.Fprintf(w, "%-10s %s -> %s (%s)\n", r.State+":", r.Edition, r.Path, r.Error); err != nil {
return err
}
continue
}
if _, err := fmt.Fprintf(w, "%-10s %s -> %s (%s)\n", r.State+":", r.Edition, r.Path, r.Date); err != nil {
return err
}
}
return nil
}
}

36
net/geoip/cmd/geoip-update/main_test.go Normal file
View File

@ -0,0 +1,36 @@
package main
import (
"bytes"
"strings"
"testing"
)
func TestParseFormat(t *testing.T) {
for _, format := range []string{"pretty", "tsv", "csv", "json"} {
got, err := parseFormat(format)
if err != nil {
t.Errorf("parseFormat(%q): %v", format, err)
}
if got != format {
t.Errorf("parseFormat(%q) = %q, want %q", format, got, format)
}
}
if _, err := parseFormat("yaml"); err == nil {
t.Error("parseFormat(yaml): expected error")
}
}
func TestWriteResultsTSVHasNoHeader(t *testing.T) {
var out bytes.Buffer
results := []result{{State: "updated", Edition: "GeoLite2-City", Path: "/tmp/city.tar.gz", Date: "2026-09-08"}}
if err := writeResults(&out, "tsv", results); err != nil {
t.Fatalf("writeResults: %v", err)
}
if strings.HasPrefix(out.String(), "state\t") {
t.Fatalf("TSV has a header: %q", out.String())
}
if !strings.HasPrefix(out.String(), "updated\tGeoLite2-City\t") {
t.Fatalf("TSV row = %q", out.String())
}
}