-
Notifications
You must be signed in to change notification settings - Fork 0
/
wordlist.go
81 lines (74 loc) · 1.55 KB
/
wordlist.go
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
69
70
71
72
73
74
75
76
77
78
79
80
81
package main
import (
"bufio"
"bytes"
"fmt"
"go/format"
"io"
"io/ioutil"
"net/http"
"os"
"time"
)
type wordlist struct {
filename string
pkgname string
varname string
min int
source string
timeout time.Duration
}
func (w wordlist) generate() error {
rc, err := w.fetch()
if err != nil {
return err
}
defer rc.Close()
return w.write(rc)
}
func (w wordlist) fetch() (io.ReadCloser, error) {
f, err := os.Open(w.source)
if err != nil {
return w.fetchHTTP()
}
return f, nil
}
func (w wordlist) fetchHTTP() (io.ReadCloser, error) {
client := &http.Client{Timeout: w.timeout}
req, err := http.NewRequest(http.MethodGet, w.source, nil)
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "text/plain")
resp, err := client.Do(req)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
err = fmt.Errorf("top95: fetch status %v", resp.StatusCode)
}
return resp.Body, err
}
func (w wordlist) write(r io.Reader) error {
var buf bytes.Buffer
buf.WriteString("// generated by top95\n\n")
buf.WriteString("package " + w.pkgname + "\n")
buf.WriteString("var " + w.varname + " = map[string]struct{}{\n")
w.flush(r, &buf)
buf.WriteString("}\n")
data, err := format.Source(buf.Bytes())
if err != nil {
return err
}
return ioutil.WriteFile(w.filename, data, 0644)
}
func (w wordlist) flush(r io.Reader, buf *bytes.Buffer) error {
s := bufio.NewScanner(r)
for s.Scan() {
line := s.Text()
if len(line) > w.min {
buf.WriteString(fmt.Sprintf("%q: struct{}{},\n", line))
}
}
return s.Err()
}