AdGuardDNS/internal/geoip/asntops_generate.go
Andrey Meshkov 87137bddcf Sync v2.10.0
2024-11-08 16:26:22 +03:00

104 lines
2.6 KiB
Go

//go:build generate
package main
import (
"context"
"encoding/json"
"log/slog"
"net/http"
"os"
"text/template"
"time"
"github.com/AdguardTeam/AdGuardDNS/internal/agdhttp"
"github.com/AdguardTeam/AdGuardDNS/internal/geoip"
"github.com/AdguardTeam/golibs/errors"
"github.com/AdguardTeam/golibs/httphdr"
"github.com/AdguardTeam/golibs/logutil/slogutil"
"github.com/AdguardTeam/golibs/osutil"
)
func main() {
ctx := context.Background()
logger := slogutil.New(nil)
defer slogutil.RecoverAndExit(ctx, logger, osutil.ExitCodeFailure)
c := &http.Client{
Timeout: 10 * time.Second,
}
req := errors.Must(http.NewRequestWithContext(ctx, http.MethodGet, countriesASNURL, nil))
req.Header.Add(httphdr.UserAgent, agdhttp.UserAgent())
resp := errors.Must(c.Do(req))
defer slogutil.CloseAndLog(ctx, logger, resp.Body, slog.LevelError)
err := agdhttp.CheckStatus(resp, http.StatusOK)
errors.Check(err)
out := errors.Must(os.OpenFile("./asntops.go", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o664))
defer slogutil.CloseAndLog(ctx, logger, out, slog.LevelError)
defaultCountryTopASNs := map[geoip.Country][]geoip.ASN{}
err = json.NewDecoder(resp.Body).Decode(&defaultCountryTopASNs)
errors.Check(err)
// Don't use a *container.MapSet here, because the map is iterated over in
// the template.
defaultTopASNs := map[geoip.ASN]struct{}{}
for _, asns := range defaultCountryTopASNs {
for _, asn := range asns {
if asn != 0 {
defaultTopASNs[asn] = struct{}{}
}
}
}
type templateData struct {
DefaultTopASNs map[geoip.ASN]struct{}
DefaultCountryTopASNs map[geoip.Country][]geoip.ASN
}
tmplData := &templateData{
DefaultTopASNs: defaultTopASNs,
DefaultCountryTopASNs: defaultCountryTopASNs,
}
tmpl := template.Must(template.New("main").Parse(tmplStr))
err = tmpl.Execute(out, tmplData)
errors.Check(err)
}
// countriesASNURL is the default URL to get the per-country top ASN statistics
// from.
const countriesASNURL = `https://static.adtidy.org/dns/countries_asn.json`
// tmplStr is the template of the generated Go code.
const tmplStr = `// Code generated by go run ./asntops_generate.go; DO NOT EDIT.
package geoip
import "github.com/AdguardTeam/golibs/container"
// DefaultTopASNs contains all specially handled ASNs.
var DefaultTopASNs = container.NewMapSet[ASN](
{{- range $asn, $_ := .DefaultTopASNs }}
{{ $asn }},
{{- end }}
)
// DefaultCountryTopASNs is a mapping of a country to their top ASNs.
var DefaultCountryTopASNs = map[Country]ASN{
{{- range $ctry, $ASNs := .DefaultCountryTopASNs }}
{{- if gt (len $ASNs) 0 }}
Country{{ $ctry }}: {{ index $ASNs 0 }},
{{- else }}
{{- continue }}
{{- end }}
{{- end }}
}
`