censys.go

/*
 * ZAnnotate Copyright 2026 Regents of the University of Michigan
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package zannotate

import (
	"encoding/json"
	"flag"
	"io"
	"net"
	"net/http"
	"strings"

	log "github.com/sirupsen/logrus"
)

type CensysAnnotatorFactory struct {
	BasePluginConf
	client        *http.Client // Shared client across threads
	personalToken string       // User's personal access token for API auth
}

// Censys Annotator Factory (Global)

func (a *CensysAnnotatorFactory) MakeAnnotator(i int) Annotator {
	var v CensysAnnotator
	v.Factory = a
	v.Id = i
	return &v
}

func (a *CensysAnnotatorFactory) Initialize(_ *GlobalConf) error {
	a.client = http.DefaultClient
	return nil
}

func (a *CensysAnnotatorFactory) GetWorkers() int {
	return a.Threads
}

func (a *CensysAnnotatorFactory) Close() error {
	return nil
}

func (a *CensysAnnotatorFactory) IsEnabled() bool {
	return a.Enabled
}

func (a *CensysAnnotatorFactory) GroupName() string { return "Censys" }

func (a *CensysAnnotatorFactory) AddFlags(flags *flag.FlagSet) {
	flags.BoolVar(&a.Enabled, "censys", false, "annotate with censys internet intelligence")
	flags.StringVar(&a.personalToken, "censys-pat", "", "censys API personal access token (PAT)")
	flags.IntVar(&a.Threads, "censys-threads", 1, "how many enrichment threads to use. Note that free plan only allows 1 concurrent API request at a time")
}

// CensysAnnotator (Per-Worker)
type CensysAnnotator struct {
	Factory *CensysAnnotatorFactory
	Id      int
}

func (a *CensysAnnotator) Initialize() (err error) {
	return nil
}

func (a *CensysAnnotator) GetFieldName() string {
	return "censys"
}

var censysAPIHostLookupURL = "https://api.platform.censys.io/v3/global/asset/host/"

// Annotate performs a Censys host lookup for the given IP address and returns the results.
// If an error occurs or a lookup fails, it returns nil
func (a *CensysAnnotator) Annotate(ip net.IP) interface{} {
	req, err := http.NewRequest("GET", censysAPIHostLookupURL+ip.String(), nil)
	if err != nil {
		// If we can't even form a request, we'll fail to enrich anything. Erroring out.
		log.Fatalf("could not form an http request for enriching with censys data for ip %s: %v", ip.String(), err)
	}

	req.Header.Add("accept", "application/json")
	req.Header.Add("authorization", "Bearer "+a.Factory.personalToken)

	res, err := a.Factory.client.Do(req)
	if err != nil {
		log.Debugf("failed to annotate ip %s with censys: %v", ip.String(), err)
		return nil
	}
	defer func(Body io.ReadCloser) {
		err = Body.Close()
		if err != nil {
			log.Debugf("failed to close response body: %v", err)
		}
	}(res.Body)
	body, _ := io.ReadAll(res.Body)
	if res.StatusCode >= 400 && res.StatusCode < 500 {
		// From https://docs.censys.com/reference/get-started#step-6-handle-http-response-codes
		// 4XX errors are not transient and so we'll abort and report to the user
		log.Fatalf("censys api returned an http status '%s' with message: '%s'. "+
			"Cannot continue to annotate, please check that you have sufficient API credits and your PAT is correct", res.Status, strings.TrimSpace(string(body)))

	} else if res.StatusCode != http.StatusOK {
		// Should be a transient error, log and move on
		log.Debugf("censys api returned an http %s status with message: '%s'. Skipping Censys annotation for this IP: %s", res.Status, strings.TrimSpace(string(body)), ip.String())
		return nil
	}
	// We have a successful response, unmarshall it.
	// Struct taken from v1.1 of Censys API docs
	var result struct {
		Result struct {
			Resource any `json:"resource"`
		} `json:"result"`
	}
	err = json.Unmarshal(body, &result)
	if err != nil {
		log.Debugf("failed to parse censys response for ip %s: %v", ip.String(), err)
		return nil
	}
	return result.Result.Resource
}

func (a *CensysAnnotator) Close() error {
	return nil
}

func init() {
	s := new(CensysAnnotatorFactory)
	RegisterAnnotator(s)
}