Test whether given IP belongs to AWS network

main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/netip"
	"os"
	"strings"
)

func main() {
	log.SetFlags(0)
	flag.Parse()
	if err := run(context.Background(), flag.Arg(0)); err != nil {
		log.Fatal(err)
	}
}

func run(ctx context.Context, arg string) error {
	if arg == "" {
		return errors.New("want a single ip as the first argument")
	}
	addr, err := netip.ParseAddr(arg)
	if err != nil {
		return err
	}
	prefixes, err := loadPrefixes(ctx)
	if err != nil {
		return err
	}
	var matched bool
	for _, p := range prefixes {
		if !p.prefix.Contains(addr) {
			continue
		}
		fmt.Printf("%v belongs to %s, %s (%s)\n", addr, p.prefix, p.service, p.region)
		matched = true
	}
	if matched {
		return nil
	}
	return fmt.Errorf("%v does not belong to any of %d prefixes", addr, len(prefixes))
}

func loadPrefixes(ctx context.Context) ([]awsPrefix, error) {
	decode := func(r io.Reader) ([]awsPrefix, error) {
		var tmp struct {
			Prefixes []awsPrefix `json:"prefixes"`
		}
		if err := json.NewDecoder(r).Decode(&tmp); err != nil {
			return nil, err
		}
		if len(tmp.Prefixes) == 0 {
			return nil, errors.New("no prefixes")
		}
		return tmp.Prefixes, nil
	}
	const cacheFile = "ip-ranges.json"
	if f, err := os.Open(cacheFile); err == nil {
		defer f.Close()
		return decode(f)
	}
	// https://docs.aws.amazon.com/vpc/latest/userguide/aws-ip-ranges.html
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://ip-ranges.amazonaws.com/ip-ranges.json", nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status: %s", resp.Status)
	}
	if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "application/json") {
		return nil, fmt.Errorf("unsupported Content-Type: %s", ct)
	}
	var buf bytes.Buffer
	p, err := decode(io.TeeReader(resp.Body, &buf))
	if err != nil {
		return nil, err
	}
	return p, os.WriteFile(cacheFile, buf.Bytes(), 0666)
}

type awsPrefix struct {
	prefix  netip.Prefix
	region  string
	service string
}

func (a *awsPrefix) UnmarshalJSON(b []byte) error {
	var tmp struct {
		Prefix  string `json:"ip_prefix"`
		Region  string `json:"region"`
		Service string `json:"service"`
	}
	if err := json.Unmarshal(b, &tmp); err != nil {
		return err
	}
	p, err := netip.ParsePrefix(tmp.Prefix)
	if err != nil {
		return err
	}
	a.prefix = p
	a.region = tmp.Region
	a.service = tmp.Service
	return nil
}