forked from CovenantSQL/CookieScanner
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
115 lines (99 loc) · 2.91 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
114
115
/*
* Copyright 2019 The CovenantSQL Authors.
*
* 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 cli
import (
"fmt"
"os"
"github.com/CovenantSQL/CookieScanner/cmd"
"github.com/CovenantSQL/CookieScanner/parser"
"github.com/pkg/errors"
kingpin "gopkg.in/alecthomas/kingpin.v2"
)
var (
port int
headless bool
outputJSON bool
outputHTML string
outputPDF string
site string
)
func RegisterCommand(app *kingpin.Application, opts *cmd.CommonOptions) {
c := app.Command("cli", "generate report for a single website")
c.Flag("headless", "run chrome in headless mode").BoolVar(&headless)
c.Flag("port", "chrome remote debugger listen port").Default("9222").IntVar(&port)
c.Flag("json", "print report as json").BoolVar(&outputJSON)
c.Flag("html", "save report as html").StringVar(&outputHTML)
c.Flag("pdf", "save report as pdf").StringVar(&outputPDF)
c.Arg("site", "site url").Required().StringVar(&site)
c.Action(func(context *kingpin.ParseContext) error {
return handler(opts)
})
}
func handler(opts *cmd.CommonOptions) (err error) {
if !outputJSON && outputHTML == "" && outputPDF == "" {
outputJSON = true
}
if outputPDF != "" {
headless = true
}
t := parser.NewTask(&parser.TaskConfig{
Timeout: opts.Timeout,
WaitAfterPageLoad: opts.WaitAfterPageLoad,
Verbose: opts.Verbose,
ChromeApp: opts.ChromeApp,
DebuggerPort: port,
Headless: headless,
Classifier: opts.ClassifierHandler,
})
if err = t.Start(); err != nil {
err = errors.Wrapf(err, "start debugger failed")
return
}
defer t.Cleanup()
if err = t.Parse(site); err != nil {
err = errors.Wrapf(err, "get site cookie info failed")
return
}
if outputJSON {
if jsonData, err := t.OutputJSON(true); err == nil {
fmt.Println(jsonData)
} else {
err = errors.Wrapf(err, "generate json report failed")
}
return
}
if outputHTML != "" {
var htmlData string
if htmlData, err = t.OutputHTML(); err != nil {
err = errors.Wrapf(err, "generate html report failed")
} else {
var f *os.File
if f, err = os.OpenFile(outputHTML, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755); err != nil {
err = errors.Wrap(err, "write html report failed")
return
}
_, _ = f.WriteString(htmlData)
_ = f.Sync()
_ = f.Close()
}
return
}
if outputPDF != "" {
err = errors.Wrapf(t.OutputPDFToFile(outputPDF), "generate pdf report failed")
return
}
return
}