Merge pull request #22 from piaolin/dev

pref: update output

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

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,20 +133,17 @@ 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()
 }
 
 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)

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)

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]
+	}
+}

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
+}

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
+	}
+}

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()