Merge pull request #62 from zmap/phillip/bug-53-support-csv-input

Support CSV input format

Author: zakird

Makefile

@@ -23,8 +23,15 @@ uninstall:
 	else \
 		echo "Uninstallation cancelled."; \
 	fi
-test:
+test: zannotate
 	go test -v ./...
+	@if [ ! -d ".venv" ]; then \
+		python3 -m venv .venv; \
+	fi
+	@. .venv/bin/activate && \
+		pip install -r requirements.txt -q && \
+		pytest test_zannotate.py -v
+
 
 lint:
 	goimports -w -local "github.com/zmap/zannotate" ./

README.md

@@ -164,6 +164,31 @@ echo "127.0.0.1" | zannotate --rdns --geoasn --geoasn-database=/path-to-geo-asn.
 ```
 
 ## Input
+
+### New-line Separated IPs
+By default, ZAnnotate expects new-line delimited IP addresses on standard input. For example:
+```shell
+printf "1.1.1.1\n8.8.8.8" | zannotate --rdns
+```
+
+```jsonl
+{"ip":"1.1.1.1","rdns":{"domain_names":["one.one.one.one"]}}
+{"ip":"8.8.8.8","rdns":{"domain_names":["dns.google"]}}
+```
+
+### JSON and CSV Flags
+
+The `--output-annotation-field` flag can be used to specify a different field name for the annotations instead of `zannotate` for both CSV and JSON. For example:
+
+```shell
+printf "name,ip_address,date\n cloudflare,1.1.1.1,04-04-26\n google,8.8.8.8,04-04-26" | ./zannotate --rdns --input-file-type=csv --input-ip-field=ip_address --output-annotation-field="info"
+```
+
+```json lines
+{"name":" cloudflare","ip_address":"1.1.1.1","date":"04-04-26","info":{"rdns":{"domain_names":["one.one.one.one"]}}}
+{"ip_address":"8.8.8.8","date":"04-04-26","info":{"rdns":{"domain_names":["dns.google"]}},"name":" google"}
+```
+### JSON
 You may wish to annotate data that is already in JSON format. You'll then need to use the `--input-file-type=json` flag.
 This will insert a `zannotate` field into the existing JSON object. For example:
 
@@ -175,6 +200,41 @@ echo '{"ip": "1.1.1.1"}'  | ./zannotate --rdns --geoasn --geoasn-database=/path-
 {"ip":"1.1.1.1","zannotate":{"geoasn":{"asn":13335,"org":"CLOUDFLARENET"},"rdns":{"domain_names":["one.one.one.one"]}}}
 ```
 
+If your JSON objects have a different field for the IP address, you can specify that with the `--input-ip-field` flag. For example, if your JSON objects have an `ip_address` field instead of `ip`, you can use:
+
+```shell
+echo '{"ip_address": "1.1.1.1"}'  | ./zannotate --rdns --input-file-type=json --input-ip-field=ip_address    
+```
+
+```json
+{"ip_address":"1.1.1.1","zannotate":{"rdns":{"domain_names":["one.one.one.one"]}}}
+````
+
+### CSV
+If your input data is in CSV format, you can use the `--input-file-type=csv` flag.
+
+```shell
+printf "name,ip,date\n cloudflare,1.1.1.1,04-04-26\n google,8.8.8.8,04-04-26" | ./zannotate --rdns --input-file-type=csv 
+```
+
+```jsonl
+{"name":" cloudflare","ip":"1.1.1.1","date":"04-04-26","zannotate":{"rdns":{"domain_names":["one.one.one.one"]}}}
+{"name":" google","ip":"8.8.8.8","date":"04-04-26","zannotate":{"rdns":{"domain_names":["dns.google"]}}}
+```
+
+Similar to JSON, you can use the `--input-ip-field` flag to specify a column other than `ip` that contains the IP address.
+
+```shell
+printf "name,ip_address,date\n cloudflare,1.1.1.1,04-04-26\n google,8.8.8.8,04-04-26" | ./zannotate --rdns --input-file-type=csv --input-ip-field=ip_address
+```
+
+```jsonl
+{"date":"04-04-26","zannotate":{"rdns":{"domain_names":["dns.google"]}},"name":" google","ip_address":"8.8.8.8"}
+{"date":"04-04-26","zannotate":{"rdns":{"domain_names":["one.one.one.one"]}},"name":" cloudflare","ip_address":"1.1.1.1"}
+```
+
+
+
 # Modules
 
 ## RDAP (WHOIS successor)

annotate.go

@@ -16,6 +16,8 @@ package zannotate
 
 import (
 	"bufio"
+	"encoding/csv"
+	"slices"
 
 	"io"
 	"net"
@@ -68,6 +70,41 @@ func ipToInProcess(line string) inProcessIP {
 	return retv
 }
 
+func csvToInProcess(record []string, headers []string, ipFieldName, annotationFieldName string) inProcessIP {
+	var retv inProcessIP
+
+	if len(record) != len(headers) {
+		log.Fatalf("csv record length (%d) does not match header length (%d) for record: %v", len(record), len(headers), record)
+	}
+
+	outMap := make(map[string]interface{}, len(headers))
+	for i, header := range headers {
+		outMap[header] = record[i]
+	}
+
+	foundIPField := false
+	for i, header := range headers {
+		switch header {
+		case ipFieldName:
+			foundIPField = true
+			retv.Ip = net.ParseIP(record[i])
+			if retv.Ip == nil {
+				log.Fatalf("unable to parse IP at column '%d': %s", i, record[i])
+			}
+		case annotationFieldName:
+			log.Fatalf("csv headers already contains annotation key '%v'", annotationFieldName)
+		default:
+			outMap[header] = record[i]
+		}
+	}
+	if !foundIPField {
+		log.Fatalf("unable to find IP address field with IP field name %s in CSV record: %v", ipFieldName, record)
+	}
+
+	retv.Out = outMap
+	return retv
+}
+
 // single worker that reads from file and queues raw lines
 func AnnotateRead(conf *GlobalConf, path string, in chan<- string) {
 	log.Debug("read thread started")
@@ -84,7 +121,20 @@ func AnnotateRead(conf *GlobalConf, path string, in chan<- string) {
 		log.Debug("reading input from ", path)
 	}
 	r := bufio.NewReader(f)
-	// read IPs out of JSON input
+	if conf.InputFileType == "csv" {
+		// Need to extract the header
+		header, err := r.ReadString('\n')
+		if err != nil {
+			log.Fatal("unable to read the CSV headers", err.Error())
+		}
+		csvReader := csv.NewReader(strings.NewReader(header))
+		fields, err := csvReader.Read()
+		if err != nil {
+			log.Fatal("unable to parse CSV headers: ", err.Error())
+		}
+		conf.csvHeaders = fields
+	}
+	// read IPs out of input
 	for {
 		line, err := r.ReadString('\n')
 		if line != "" {
@@ -107,14 +157,28 @@ func AnnotateInputDecode(conf *GlobalConf, inChan <-chan string,
 	outChan chan<- inProcessIP, wg *sync.WaitGroup, i int) {
 	for line := range inChan {
 		l := strings.TrimSuffix(line, "\n")
-		if conf.InputFileType == "json" {
-			val := jsonToInProcess(l, conf.JSONIPFieldName, conf.JSONAnnotationFieldName)
-			if conf.JSONAnnotationFieldName != "" {
-				val.Out[conf.JSONAnnotationFieldName] = make(map[string]interface{})
+		switch conf.InputFileType {
+		case "json":
+			val := jsonToInProcess(l, conf.InputIPFieldName, conf.OutputAnnotationFieldName)
+			if conf.OutputAnnotationFieldName != "" {
+				val.Out[conf.OutputAnnotationFieldName] = make(map[string]interface{})
 			}
 			outChan <- val
-		} else {
+		case "csv":
+			r := csv.NewReader(strings.NewReader(l))
+			row, err := r.Read()
+			if err != nil {
+				log.Errorf("failed to parse CSV line (%s): %v", l, err)
+				continue
+			}
+			val := csvToInProcess(row, conf.csvHeaders, conf.InputIPFieldName, conf.OutputAnnotationFieldName)
+			if conf.OutputAnnotationFieldName != "" {
+				val.Out[conf.OutputAnnotationFieldName] = make(map[string]interface{})
+			}
+			outChan <- val
+		default:
 			outChan <- ipToInProcess(l)
+
 		}
 	}
 	log.Debug("decode thread ", i, " done")
@@ -160,7 +224,7 @@ func AnnotateWrite(path string, out <-chan string, wg *sync.WaitGroup) {
 	log.Debug("write thread finished")
 }
 
-func AnnotateWorker(a Annotator, inChan <-chan inProcessIP,
+func AnnotateWorker(conf *GlobalConf, a Annotator, inChan <-chan inProcessIP,
 	outChan chan<- inProcessIP, fieldName string, wg *sync.WaitGroup, i int) {
 	name := a.GetFieldName()
 	log.Debug("annotate worker (", name, ") ", i, " started")
@@ -169,7 +233,7 @@ func AnnotateWorker(a Annotator, inChan <-chan inProcessIP,
 		log.Fatal("error initializing annotate worker: ", err)
 	}
 	for inProcess := range inChan {
-		if fieldName != "" {
+		if fieldName != "" && slices.Contains([]string{"json", "csv"}, conf.InputFileType) {
 			p := inProcess.Out[fieldName].(map[string]interface{})
 			p[name] = a.Annotate(inProcess.Ip)
 		} else {
@@ -198,6 +262,7 @@ func DoAnnotation(conf *GlobalConf) {
 	inDecoded := make(chan inProcessIP)
 	// read input file
 	go AnnotateRead(conf, conf.InputFilePath, inRaw)
+
 	// decode input data
 	var decodeWG sync.WaitGroup
 	for i := 0; i < conf.InputDecodeThreads; i++ {
@@ -213,8 +278,8 @@ func DoAnnotation(conf *GlobalConf) {
 		if annotator.IsEnabled() {
 			var annotateWG sync.WaitGroup
 			for i := 0; i < annotator.GetWorkers(); i++ {
-				go AnnotateWorker(annotator.MakeAnnotator(i), lastChannel, nextChannel,
-					conf.JSONAnnotationFieldName, &annotateWG, i)
+				go AnnotateWorker(conf, annotator.MakeAnnotator(i), lastChannel, nextChannel,
+					conf.OutputAnnotationFieldName, &annotateWG, i)
 				annotateWG.Add(1)
 			}
 			lastChannel = nextChannel

cmd/zannotate/main.go

@@ -17,6 +17,7 @@ package main
 import (
 	"flag"
 	"os"
+	"slices"
 
 	log "github.com/sirupsen/logrus"
 
@@ -35,8 +36,8 @@ func main() {
 	flags.StringVar(&conf.LogFilePath, "log-file", "", "where should JSON logs be saved")
 	flags.IntVar(&conf.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
 	// json annotation configuration
-	flags.StringVar(&conf.JSONIPFieldName, "json-ip-field", "ip", "key in JSON that contains IP address")
-	flags.StringVar(&conf.JSONAnnotationFieldName, "json-annotation-field", "zannotate", "key that metadata is injected at")
+	flags.StringVar(&conf.InputIPFieldName, "input-ip-field", "ip", "key in JSON or column in CSV that contains IP address")
+	flags.StringVar(&conf.OutputAnnotationFieldName, "output-annotation-field", "zannotate", "key that metadata is injected at, used for both CSV and JSON file inputs to preserve data in the input file")
 	// encode/decode threads
 	flags.IntVar(&conf.InputDecodeThreads, "input-decode-threads", 3, "number of golang processes to decode input data (e.g., json)")
 	flags.IntVar(&conf.OutputEncodeThreads, "output-encode-threads", 3, "number of golang processes to encode output data (e.g., json)")
@@ -71,7 +72,7 @@ func main() {
 	default:
 		log.Fatal("Unknown verbosity level specified. Must be between 1 (lowest)--5 (highest)")
 	}
-	if conf.InputFileType != "ips" && conf.InputFileType != "json" {
+	if !slices.Contains([]string{"ips", "json", "csv"}, conf.InputFileType) {
 		log.Fatal("invalid input file type")
 	}
 	// check if we have any annotations to be performed
@@ -87,7 +88,7 @@ func main() {
 	}
 	// perform sanity checks
 	if conf.InputFileType == "ips" {
-		conf.JSONAnnotationFieldName = ""
+		conf.OutputAnnotationFieldName = ""
 	}
 
 	// perform annotation

conf.go

@@ -45,17 +45,19 @@ type BasePluginConf struct {
 // global library configuration
 
 type GlobalConf struct {
-	InputFilePath           string
-	InputFileType           string
-	OutputFilePath          string
-	MetadataFilePath        string
-	LogFilePath             string
-	Verbosity               int
-	Threads                 int
-	JSONIPFieldName         string
-	JSONAnnotationFieldName string
-	InputDecodeThreads      int
-	OutputEncodeThreads     int
+	InputFilePath             string
+	InputFileType             string
+	OutputFilePath            string
+	MetadataFilePath          string
+	LogFilePath               string
+	Verbosity                 int
+	Threads                   int
+	InputIPFieldName          string
+	OutputAnnotationFieldName string
+	InputDecodeThreads        int
+	OutputEncodeThreads       int
+
+	csvHeaders []string // Header row of a CSV file, used when parsing a CSV input
 }
 
 var Annotators []AnnotatorFactory

requirements.txt

@@ -0,0 +1,5 @@
+iniconfig==2.3.0
+packaging==26.0
+pluggy==1.6.0
+Pygments==2.20.0
+pytest==9.0.2