anubis/internal/thoth/thothmock/iptoasn.go

package thothmock

import (
	"context"
	"net/netip"

	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func MockIpToASNService() *IpToASNService {
	responses := map[string]*iptoasnv1.LookupResponse{
		"127.0.0.1": {Announced: false},
		"::1":       {Announced: false},
		"10.10.10.10": {
			Announced:   true,
			AsNumber:    13335,
			Cidr:        []string{"1.1.1.0/24"},
			CountryCode: "US",
			Description: "Cloudflare",
		},
		"2.2.2.2": {
			Announced:   true,
			AsNumber:    420,
			Cidr:        []string{"2.2.2.0/24"},
			CountryCode: "CA",
			Description: "test canada",
		},
		"1.1.1.1": {
			Announced:   true,
			AsNumber:    13335,
			Cidr:        []string{"1.1.1.0/24"},
			CountryCode: "US",
			Description: "Cloudflare",
		},
	}

	return &IpToASNService{Responses: responses}
}

type IpToASNService struct {
	iptoasnv1.UnimplementedIpToASNServiceServer
	Responses map[string]*iptoasnv1.LookupResponse
}

func (ip2asn *IpToASNService) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) {
	if _, err := netip.ParseAddr(lr.GetIpAddress()); err != nil {
		return nil, err
	}

	resp, ok := ip2asn.Responses[lr.GetIpAddress()]
	if !ok {
		return nil, status.Error(codes.NotFound, "IP address not found in mock")
	}

	return resp, nil
}