File: csv.go

package info (click to toggle)
hellfire 0.0~git20180708.bf3c390-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, buster, sid
  • size: 96 kB
  • sloc: makefile: 2
file content (68 lines) | stat: -rw-r--r-- 1,382 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
package hellfire // import "pathspider.net/hellfire"

import (
	"bufio"
	"encoding/csv"
	"io"
	"net/url"
	"os"
)

// A CSVList handles input in CSV format. There may be a more specific type
// available and that should be used if that is the case (e.g. for the
// Alexa or Citizen Lab test lists).
type CSVList struct {
	TestList
	reader io.Reader
	header []string
}

func CSVListFromFile(filename string) *CSVList {
	f, err := os.Open(filename)
	if err != nil {
		panic("Error opening file")
	}
	return CSVListFromReader(f)
}

func CSVListFromReader(reader io.Reader) *CSVList {
	l := new(CSVList)
	l.reader = reader
	return l
}

func (l *CSVList) SetHeader(header []string) {
	l.header = header
}

func (l *CSVList) FeedJobs(jobs chan map[string]interface{}) {
	reader := csv.NewReader(bufio.NewReader(l.reader))
	if reader == nil {
		panic("CSVList not initialised with a reader")
	}
	var header []string
	if l.header == nil {
		var err error
		header, err = reader.Read()
		if err != nil {
			panic("Error reading the header from the CSV")
		}
	} else {
		header = l.header
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		r := make(map[string]interface{})
		for idx, name := range header {
			r[name] = record[idx]
		}
		if r["domain"] == nil && r["url"] != nil {
			u, _ := url.Parse(r["url"].(string))
			r["domain"] = u.Host
		}
		jobs <- r
	}
}