From a67d2dc3ad3e641df0594e7b60d10d48f812dd5f Mon Sep 17 00:00:00 2001 From: lxy <1418676300@qq.com> Date: Mon, 22 May 2023 04:18:21 -0700 Subject: [PATCH] pref: update output Made improvement - extraction function - update output - fix screenshot - fix NSFW options - fix update command --- cmd/detect.go | 5 +-- cmd/screenshot.go | 88 ++++++++++----------------------------- cmd/update.go | 23 ++++------ utils/addTarget.go | 7 ++++ utils/parseResult.go | 25 +++++++++++ utils/renameFileByTime.go | 32 ++++++++++++++ utils/writeToFile.go | 7 ++++ 7 files changed, 105 insertions(+), 82 deletions(-) create mode 100644 utils/addTarget.go create mode 100644 utils/parseResult.go create mode 100644 utils/renameFileByTime.go diff --git a/cmd/detect.go b/cmd/detect.go index 74f3571..c203e27 100644 --- a/cmd/detect.go +++ b/cmd/detect.go @@ -36,7 +36,7 @@ var ( wg sync.WaitGroup nonSiteData = "[-] There is no site data of %s\n" existInfo = "[+] %-15s %-15s: %s\n" - existOutputInfo = "%-15s %-15s: %s\n" + existOutputInfo = "%s, %s, %s\n" nonExistInfo = "[-] %-15s %-15s: non exists\n" reqErrorInfo = "[!] %-15s %-15s: %s requests error, retry %d/%d\n" sleepMap = make(map[string]int64) @@ -163,7 +163,7 @@ func detectSite(name, site, nameType string, siteBody gjson.Result) { // flag for precisely mode flag := false - if siteBody.Get("isNFSW").Bool() && !detectArgs.isNSFW { + if siteBody.Get("isNSFW").Bool() && !detectArgs.isNSFW { log.Debugf(nsfwInfo, name, site, site) return } @@ -226,7 +226,6 @@ func detectUser(name, site string, requestTimes, retryTimes, detectCount int, fl strValue := fmt.Sprintf("%v", value) header[strKey] = strValue } - } // if body is set, set method to post diff --git a/cmd/screenshot.go b/cmd/screenshot.go index 00a072b..c096dc7 100644 --- a/cmd/screenshot.go +++ b/cmd/screenshot.go @@ -1,11 +1,10 @@ package cmd import ( - "bufio" + "DetectDee/utils" "context" "fmt" "io/ioutil" - "net/url" "os" "strings" "sync" @@ -29,9 +28,8 @@ type screenshotArgsType struct { var ( screenshotArgs screenshotArgsType screenshotWg sync.WaitGroup - targets = make(chan []string, 10) + screenshotTargets = make(chan []string, 10) execPathNotFound = "[-] Chrome executable file not found, please install Chrome or specify the chrome.exe path with --path" - urlFileReadError = "[-] failed to read targets file:%v" createResultFolderError = "[-] failed to create result folder:%v" ) @@ -59,9 +57,14 @@ func screenshot(_ *cobra.Command, _ []string) { log.SetLevel(log.DebugLevel) } - go parseFile(screenshotArgs.file) + targetList, err := utils.ParseResult(screenshotArgs.file) + if err != nil { + return + } - err := CreateDirIfNotExists(screenshotArgs.result) + go utils.AddTarget(targetList, screenshotTargets) + + err = CreateDirIfNotExists(screenshotArgs.result) if err != nil { log.Errorf(createResultFolderError, err) return @@ -75,42 +78,11 @@ func screenshot(_ *cobra.Command, _ []string) { screenshotWg.Wait() } -func parseFile(file string) { - f, err := os.Open(file) - if err != nil { - log.Errorf(urlFileReadError, err) - return - } - defer f.Close() - - scanner := bufio.NewScanner(f) - for scanner.Scan() { - if scanner.Text() != "" { - tmp := strings.Split(scanner.Text(), ": ") - name := strings.Split(strings.TrimSpace(tmp[0]), " ") - - target := []string{fmt.Sprintf("%s-%s", name[0], name[len(name)-1]), tmp[1]} - targets <- target - } - } -} - func navigate(workerNum int) { defer screenshotWg.Done() // create context - opts := append(chromedp.DefaultExecAllocatorOptions[:], - chromedp.Flag("headless", !screenshotArgs.chrome), - chromedp.ProxyServer(screenshotArgs.proxy), - chromedp.Flag("mute-audio", true), - chromedp.IgnoreCertErrors, - chromedp.DisableGPU, - chromedp.NoFirstRun, - chromedp.ExecPath(screenshotArgs.path), - chromedp.WindowSize(1920, 1080), - chromedp.NoDefaultBrowserCheck, - chromedp.NoSandbox, - ) + opts := append(chromedp.DefaultExecAllocatorOptions[:], chromedp.Flag("headless", !screenshotArgs.chrome), chromedp.ProxyServer(screenshotArgs.proxy), chromedp.Flag("mute-audio", true), chromedp.IgnoreCertErrors, chromedp.DisableGPU, chromedp.NoFirstRun, chromedp.ExecPath(screenshotArgs.path), chromedp.WindowSize(1920, 1080), chromedp.NoDefaultBrowserCheck, chromedp.NoSandbox) if screenshotArgs.proxy != "" { opts = append(opts, chromedp.Flag("proxy-server", screenshotArgs.proxy)) @@ -123,12 +95,9 @@ func navigate(workerNum int) { ctx, cancel = chromedp.NewContext(ctx) // Check & Make Browser not close - err := chromedp.Run( - ctx, - chromedp.Tasks{ - chromedp.Navigate("https://github.com/piaolin"), - }, - ) + err := chromedp.Run(ctx, chromedp.Tasks{ + chromedp.Navigate("https://github.com/piaolin"), + }) if err != nil { if strings.Contains(err.Error(), "executable file not found") { log.Fatalln(execPathNotFound) @@ -138,12 +107,13 @@ func navigate(workerNum int) { } // Run - for len(targets) > 0 { - target := <-targets + for len(screenshotTargets) > 0 { + target := <-screenshotTargets name := target[0] - urlStr := target[1] + site := target[1] + urlStr := target[2] - log.Debugf("[Worker%2d] starting screenshot task %s,remaining tasks:%d\n", workerNum, urlStr, len(targets)) + log.Debugf("[Worker%2d] starting screenshot task %s,remaining tasks:%d\n", workerNum, urlStr, len(screenshotTargets)) var buf []byte //ctx, cancel = chromedp.NewContext( // ctx, @@ -152,7 +122,7 @@ func navigate(workerNum int) { log.Errorf("[-] Failed to take (URL:%s) screenshot: %v", urlStr, err) continue } - pngFilePath := fmt.Sprintf("./%s/%s.png", screenshotArgs.result, name) + pngFilePath := fmt.Sprintf("./%s/%s-%s.png", screenshotArgs.result, name, site) if err := ioutil.WriteFile(pngFilePath, buf, 0644); err != nil { log.Errorf("[-] Failed to write file %v", err) continue @@ -163,7 +133,7 @@ func navigate(workerNum int) { //} log.Infof("[+] screenshot success %s", urlStr) - log.Debugf("[Worker%2d] finished screenshot task %s,remaining tasks:%d\n", workerNum, urlStr, len(targets)) + log.Debugf("[Worker%2d] finished screenshot task %s,remaining tasks:%d\n", workerNum, urlStr, len(screenshotTargets)) } defer cancel() @@ -171,12 +141,9 @@ func navigate(workerNum int) { func fullScreenshot(urlstr string, quality int, res *[]byte) chromedp.Tasks { return chromedp.Tasks{ - chromedp.Navigate(urlstr), - //chromedp.WaitVisible("style"), - chromedp.Sleep(1 * time.Second), - //chromedp.OuterHTML(`document.querySelector("body")`, &htmlContent, chromedp.ByJSPath), - chromedp.FullScreenshot(res, quality), - //chromedp.ActionFunc(func(ctx context.Context) error { + chromedp.Navigate(urlstr), //chromedp.WaitVisible("style"), + chromedp.Sleep(1 * time.Second), //chromedp.OuterHTML(`document.querySelector("body")`, &htmlContent, chromedp.ByJSPath), + chromedp.FullScreenshot(res, quality), //chromedp.ActionFunc(func(ctx context.Context) error { // _, _, _, _, _, contentSize, err := page.GetLayoutMetrics().Do(ctx) // if err != nil { // return err @@ -211,15 +178,6 @@ func fullScreenshot(urlstr string, quality int, res *[]byte) chromedp.Tasks { } } -func getDomain(urlstr string) string { - u, err := url.Parse(urlstr) - if err != nil { - fmt.Println(err) - return "getdomain_error" - } - return u.Hostname() -} - func CreateDirIfNotExists(dirName string) error { if _, err := os.Stat(dirName); os.IsNotExist(err) { err := os.MkdirAll(dirName, 0755) diff --git a/cmd/update.go b/cmd/update.go index 4fdac23..7c64ffb 100644 --- a/cmd/update.go +++ b/cmd/update.go @@ -2,12 +2,8 @@ package cmd import ( "DetectDee/utils" - "fmt" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" - "os" - "strings" - "time" ) type updateArgsType struct { @@ -40,19 +36,18 @@ func update(_ *cobra.Command, _ []string) { log.SetLevel(log.DebugLevel) } - if _, err := os.Stat(updateArgs.file); err == nil { - newFileName := strings.Split(updateArgs.file, ".")[0] + "_" + strings.Replace(time.Now().Format("15:04:05"), ":", "", -1) + ".json" - log.Debugln(newFileName) - err := os.Rename(updateArgs.file, newFileName) - if err != nil { - log.Fatalln(err) - } + newFileName, err := utils.RenameFileByTime(updateArgs.file) + if err != nil { + log.Fatal(err) } - err := utils.DownloadFile(updateArgs.file, updateArgs.url) + err = utils.DownloadFile(updateArgs.file, updateArgs.url) if err != nil { - fmt.Println("Error downloading file: ", err) - return + // if download error, restore filename + if newFileName != updateArgs.file { + _ = utils.RestoreFilename(newFileName) + } + log.Fatalln("Error downloading file: ", err) } log.Infof("[+] download file to %s\n", updateArgs.file) diff --git a/utils/addTarget.go b/utils/addTarget.go new file mode 100644 index 0000000..c506a09 --- /dev/null +++ b/utils/addTarget.go @@ -0,0 +1,7 @@ +package utils + +func AddTarget(targetList [][]string, targets chan []string) { + for i := 0; i < len(targetList); i++ { + targets <- targetList[i] + } +} diff --git a/utils/parseResult.go b/utils/parseResult.go new file mode 100644 index 0000000..a559d03 --- /dev/null +++ b/utils/parseResult.go @@ -0,0 +1,25 @@ +package utils + +import ( + "bufio" + "os" + "strings" +) + +func ParseResult(filename string) ([][]string, error) { + f, err := os.Open(filename) + if err != nil { + return nil, err + } + defer f.Close() + + var resultList [][]string + scanner := bufio.NewScanner(f) + for scanner.Scan() { + if scanner.Text() != "" { + splitStr := strings.Split(scanner.Text(), ", ") + resultList = append(resultList, splitStr) + } + } + return resultList, err +} diff --git a/utils/renameFileByTime.go b/utils/renameFileByTime.go new file mode 100644 index 0000000..d9f1afa --- /dev/null +++ b/utils/renameFileByTime.go @@ -0,0 +1,32 @@ +package utils + +import ( + log "github.com/sirupsen/logrus" + "os" + "strings" + "time" +) + +func RenameFileByTime(filename string) (string, error) { + if _, err := os.Stat(filename); err == nil { + splitStr := strings.Split(filename, ".") + newFileName := splitStr[0] + "_" + time.Now().Format("20060102_150405") + "." + splitStr[1] + log.Debugln(newFileName) + err := os.Rename(filename, newFileName) + return newFileName, err + } else { + return filename, nil + } +} + +func RestoreFilename(filename string) error { + if _, err := os.Stat(filename); err == nil { + splitStr := strings.Split(filename, "_") + fileExt := strings.Split(filename, ".") + newFileName := splitStr[0] + "." + fileExt[len(fileExt)-1] + log.Debugln(newFileName) + return os.Rename(filename, newFileName) + } else { + return nil + } +} diff --git a/utils/writeToFile.go b/utils/writeToFile.go index 0420619..6fa3256 100644 --- a/utils/writeToFile.go +++ b/utils/writeToFile.go @@ -2,10 +2,17 @@ package utils import ( "fmt" + log "github.com/sirupsen/logrus" "os" ) func WriteToFile(filename string, content chan string, done chan bool) { + + _, err := RenameFileByTime(filename) + if err != nil { + log.Fatal(err) + } + file, _ := os.Create(filename) defer file.Close()