forked from google/webrisk
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
113 lines (102 loc) · 3.31 KB
/
main.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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
// Copyright 2019 Google LLC
//
// 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
//
// https://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.
// Command wrlookup is a tool for looking up URLs via the command-line.
//
// The tool reads one URL per line from STDIN and checks every URL against
// the Web Risk API. The "Safe" or "Unsafe" verdict is printed to STDOUT.
// If an error occurred, debug information may be printed to STDERR.
//
// To build the tool:
// $ go get github.com/google/webrisk/cmd/wrlookup
//
// Example usage:
// $ wrlookup -apikey $APIKEY
// https://google.com
// Safe URL: https://google.com
// http://bad1url.org
// Unsafe URL: [{bad1url.org {MALWARE ANY_PLATFORM URL}}]
package main
import (
"bufio"
"flag"
"fmt"
"github.com/google/webrisk"
"os"
)
var (
apiKeyFlag = flag.String("apikey", "", "specify your Web Risk API key")
databaseFlag = flag.String("db", "", "path to the Web Risk database. By default persistent storage is disabled (not recommended).")
serverURLFlag = flag.String("server", webrisk.DefaultServerURL, "Web Risk API server address.")
proxyFlag = flag.String("proxy", "", "proxy to use to connect to the HTTP server")
)
const usage = `wrlookup: command-line tool to lookup URLs with Web Risk.
Tool reads one URL per line from STDIN and checks every URL against the
Web Risk API. The Safe or Unsafe verdict is printed to STDOUT. If an error
occurred, debug information may be printed to STDERR.
Exit codes (bitwise OR of following codes):
0 if and only if all URLs were looked up and are safe.
1 if at least one URL is not safe.
2 if at least one URL lookup failed.
4 if the input was invalid.
Usage: %s -apikey=$APIKEY
`
const (
codeSafe = (1 << iota) / 2 // Sequence of 0, 1, 2, 4, 8, etc...
codeUnsafe
codeFailed
codeInvalid
)
func main() {
flag.Usage = func() {
fmt.Fprintf(os.Stderr, usage, os.Args[0])
flag.PrintDefaults()
}
flag.Parse()
if *apiKeyFlag == "" {
fmt.Fprintln(os.Stderr, "No -apikey specified")
os.Exit(codeInvalid)
}
sb, err := webrisk.NewWebriskClient(webrisk.Config{
APIKey: *apiKeyFlag,
DBPath: *databaseFlag,
Logger: os.Stderr,
ServerURL: *serverURLFlag,
ProxyURL: *proxyFlag,
})
if err != nil {
fmt.Fprintln(os.Stderr, "Unable to initialize Web Risk client: ", err)
os.Exit(codeInvalid)
}
scanner := bufio.NewScanner(os.Stdin)
code := codeSafe
for scanner.Scan() {
url := scanner.Text()
threats, err := sb.LookupURLs([]string{url})
if err != nil {
fmt.Fprintln(os.Stdout, "Unknown URL:", url)
fmt.Fprintln(os.Stderr, "Lookup error:", err)
code |= codeFailed
} else if len(threats[0]) == 0 {
fmt.Fprintln(os.Stdout, "Safe URL:", url)
} else {
fmt.Fprintln(os.Stdout, "Unsafe URL:", threats[0])
code |= codeUnsafe
}
}
if scanner.Err() != nil {
fmt.Fprintln(os.Stderr, "Unable to read input:", scanner.Err())
code |= codeInvalid
}
os.Exit(code)
}