From c77eb270a9e5627cd404d4f78fbacefa33be7227 Mon Sep 17 00:00:00 2001 From: Mario Bartolome Date: Fri, 26 Jun 2026 14:26:09 +0200 Subject: [PATCH] Fixes bug in Registries processing --- pkg/plugins/cidrs/pure_fns_test.go | 44 ++++++++++++++++++++++++++++++ pkg/plugins/cidrs/reverse_rir.go | 29 +++++++++++++++++--- pkg/runner/internal_test.go | 26 +++++++++++++++--- pkg/runner/run.go | 4 +-- 4 files changed, 93 insertions(+), 10 deletions(-) diff --git a/pkg/plugins/cidrs/pure_fns_test.go b/pkg/plugins/cidrs/pure_fns_test.go index c9003e7..5075641 100644 --- a/pkg/plugins/cidrs/pure_fns_test.go +++ b/pkg/plugins/cidrs/pure_fns_test.go @@ -1,9 +1,11 @@ package cidrs import ( + "encoding/json" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // ── splitHandles (rdap.go) ──────────────────────────────────────────────────── @@ -73,3 +75,45 @@ func TestIsLikelyRIRHandle_CaseSensitive(t *testing.T) { assert.False(t, isLikelyRIRHandle("SEC-FILING")) assert.True(t, isLikelyRIRHandle("PRAETORIAN-1")) // no blocklisted prefix } + +// ── ArinRefs UnmarshalJSON (reverse_rir.go) ───────────────────────────────── + +func TestArinRefs_UnmarshalArray(t *testing.T) { + data := `[{"@handle":"ACME-1","@name":"Acme"},{"@handle":"ACME-2","@name":"Acme Inc"}]` + var refs ArinRefs + require.NoError(t, json.Unmarshal([]byte(data), &refs)) + require.Len(t, refs, 2) + assert.Equal(t, "ACME-1", refs[0].Handle) + assert.Equal(t, "ACME-2", refs[1].Handle) +} + +func TestArinRefs_UnmarshalSingleObject(t *testing.T) { + data := `{"@handle":"PS-1576","@name":"Praetorian Security, Inc."}` + var refs ArinRefs + require.NoError(t, json.Unmarshal([]byte(data), &refs)) + require.Len(t, refs, 1) + assert.Equal(t, "PS-1576", refs[0].Handle) +} + +func TestArinRefs_UnmarshalNull(t *testing.T) { + data := `null` + var refs ArinRefs + require.NoError(t, json.Unmarshal([]byte(data), &refs)) + assert.Empty(t, refs) +} + +func TestArinOrgsResponse_SingleResult(t *testing.T) { + data := `{"orgs":{"orgRef":{"@handle":"PS-1576","@name":"Praetorian Security, Inc."}}}` + var resp ArinOrgsResponse + require.NoError(t, json.Unmarshal([]byte(data), &resp)) + require.Len(t, resp.Orgs.OrgRef, 1) + assert.Equal(t, "PS-1576", resp.Orgs.OrgRef[0].Handle) +} + +func TestArinOrgsResponse_MultipleResults(t *testing.T) { + data := `{"orgs":{"orgRef":[{"@handle":"PRAET-1","@name":"Praetorian Global"},{"@handle":"PS-1576","@name":"Praetorian Security"}]}}` + var resp ArinOrgsResponse + require.NoError(t, json.Unmarshal([]byte(data), &resp)) + require.Len(t, resp.Orgs.OrgRef, 2) + assert.Equal(t, "PRAET-1", resp.Orgs.OrgRef[0].Handle) +} diff --git a/pkg/plugins/cidrs/reverse_rir.go b/pkg/plugins/cidrs/reverse_rir.go index 6da9ff8..ef9378e 100644 --- a/pkg/plugins/cidrs/reverse_rir.go +++ b/pkg/plugins/cidrs/reverse_rir.go @@ -333,27 +333,48 @@ type ArinRef struct { Name string `json:"@name"` } +// ArinRefs handles ARIN's JSON quirk: single results are returned as a bare +// object, not a single-element array. Standard json.Unmarshal fails silently +// when unmarshaling an object into []T. +type ArinRefs []ArinRef + +func (r *ArinRefs) UnmarshalJSON(data []byte) error { + // Try array first (common case with multiple results). + var arr []ArinRef + if err := json.Unmarshal(data, &arr); err == nil { + *r = arr + return nil + } + // Fall back to single object (ARIN returns this for exactly one result). + var single ArinRef + if err := json.Unmarshal(data, &single); err != nil { + return err + } + *r = ArinRefs{single} + return nil +} + type ArinOrgsResponse struct { Orgs struct { - OrgRef []ArinRef `json:"orgRef"` + OrgRef ArinRefs `json:"orgRef"` } `json:"orgs"` } type ArinCustomersResponse struct { Customers struct { - CustomerRef []ArinRef `json:"customerRef"` + CustomerRef ArinRefs `json:"customerRef"` } `json:"customers"` } type ArinNetsResponse struct { Nets struct { - NetRef []ArinRef `json:"netRef"` + NetRef ArinRefs `json:"netRef"` } `json:"nets"` } type ArinAsnsResponse struct { Asns struct { - AsnRef []ArinRef `json:"asnRef"` + AsnRef ArinRefs `json:"asnRef"` } `json:"asns"` } diff --git a/pkg/runner/internal_test.go b/pkg/runner/internal_test.go index c5d4285..a4cb597 100644 --- a/pkg/runner/internal_test.go +++ b/pkg/runner/internal_test.go @@ -43,16 +43,14 @@ func TestEnrichWithHandles_MultipleRegistries(t *testing.T) { assert.Empty(t, result.Meta["afrinic_handles"]) } -func TestEnrichWithHandles_UnknownRegistryBroadcastsToAll(t *testing.T) { - // Tests Phase 11 fix: a FindingCIDRHandle with no/empty registry key - // must be broadcast to ALL FOUR RIRs (arin, ripe, apnic, afrinic). +func TestEnrichWithHandles_EmptyRegistryBroadcastsToAll(t *testing.T) { input := plugins.Input{OrgName: "Acme", Meta: make(map[string]string)} findings := []plugins.Finding{ { Type: plugins.FindingCIDRHandle, Value: "UNKNOWN-HANDLE", Source: "edgar", - Data: map[string]any{}, // no "registry" key → unknown + Data: map[string]any{}, // no "registry" key → empty string }, } result := enrichWithHandles(input, findings) @@ -60,6 +58,26 @@ func TestEnrichWithHandles_UnknownRegistryBroadcastsToAll(t *testing.T) { assert.Equal(t, "UNKNOWN-HANDLE", result.Meta["ripe_handles"]) assert.Equal(t, "UNKNOWN-HANDLE", result.Meta["apnic_handles"]) assert.Equal(t, "UNKNOWN-HANDLE", result.Meta["afrinic_handles"]) + assert.Equal(t, "UNKNOWN-HANDLE", result.Meta["lacnic_handles"]) +} + +func TestEnrichWithHandles_UnknownRegistryBroadcastsToAll(t *testing.T) { + input := plugins.Input{OrgName: "Acme", Meta: make(map[string]string)} + findings := []plugins.Finding{ + { + Type: plugins.FindingCIDRHandle, + Value: "EDGAR-HANDLE", + Source: "edgar", + Data: map[string]any{"registry": "unknown", "org": "Acme"}, + }, + } + result := enrichWithHandles(input, findings) + assert.Equal(t, "EDGAR-HANDLE", result.Meta["arin_handles"]) + assert.Equal(t, "EDGAR-HANDLE", result.Meta["ripe_handles"]) + assert.Equal(t, "EDGAR-HANDLE", result.Meta["apnic_handles"]) + assert.Equal(t, "EDGAR-HANDLE", result.Meta["afrinic_handles"]) + assert.Equal(t, "EDGAR-HANDLE", result.Meta["lacnic_handles"]) + assert.Empty(t, result.Meta["unknown_handles"]) } func TestEnrichWithHandles_PreservesExistingMeta(t *testing.T) { diff --git a/pkg/runner/run.go b/pkg/runner/run.go index a31cbd8..efb4cb8 100644 --- a/pkg/runner/run.go +++ b/pkg/runner/run.go @@ -357,8 +357,8 @@ func enrichWithHandles(input plugins.Input, findings []plugins.Finding) plugins. continue } reg, _ := f.Data["registry"].(string) - if reg == "" { - for _, r := range []string{"arin", "ripe", "apnic", "afrinic"} { + if reg == "" || reg == "unknown" { + for _, r := range []string{"arin", "ripe", "apnic", "afrinic", "lacnic"} { groups[r] = append(groups[r], f.Value) } continue