diff --git a/wrp.go b/wrp.go index a361c17..b18c59c 100644 --- a/wrp.go +++ b/wrp.go @@ -1,90 +1,94 @@ package main import ( - "context" - "fmt" - "log" - "net/http" - "strconv" - "strings" - "time" + "context" + "flag" + "fmt" + "log" + "net/http" + "strconv" + "strings" + "time" - "github.com/chromedp/cdproto/cdp" - "github.com/chromedp/cdproto/dom" - "github.com/chromedp/chromedp" + "github.com/chromedp/cdproto/cdp" + "github.com/chromedp/cdproto/dom" + "github.com/chromedp/chromedp" ) var ( - ctx context.Context - cancel context.CancelFunc - scrcap []byte + ctx context.Context + cancel context.CancelFunc + scrcap []byte ) func pageServer(out http.ResponseWriter, req *http.Request) { - req.ParseForm() - furl := req.Form["url"] - var url string - if len(furl) >= 1 && len(furl[0]) > 4 { - url = furl[0] - } else { - url = "https://www.google.com/" - } - log.Printf("%s Page Reqest for %s URL=%s\n", req.RemoteAddr, req.URL.Path, url) - out.Header().Set("Content-Type", "text/html") - fmt.Fprintf(out, "\n
\n") - if len(url) > 4 { - capture(url, out) - } - fmt.Fprintf(out, "\n\n") + req.ParseForm() + furl := req.Form["url"] + var url string + if len(furl) >= 1 && len(furl[0]) > 4 { + url = furl[0] + } else { + url = "https://www.google.com/" + } + log.Printf("%s Page Reqest for %s URL=%s\n", req.RemoteAddr, req.URL.Path, url) + out.Header().Set("Content-Type", "text/html") + fmt.Fprintf(out, "\n
\n") + if len(url) > 4 { + capture(url, out) + } + fmt.Fprintf(out, "\n\n") } func imgServer(out http.ResponseWriter, req *http.Request) { - log.Printf("%s Img Reqest for %s\n", req.RemoteAddr, req.URL.Path) - out.Header().Set("Content-Type", "image/png") - out.Header().Set("Content-Length", strconv.Itoa(len(scrcap))) - out.Write(scrcap) + log.Printf("%s Img Reqest for %s\n", req.RemoteAddr, req.URL.Path) + out.Header().Set("Content-Type", "image/png") + out.Header().Set("Content-Length", strconv.Itoa(len(scrcap))) + out.Write(scrcap) } func capture(url string, out http.ResponseWriter) { - var nodes []*cdp.Node - ctxx := chromedp.FromContext(ctx) - var target string + var nodes []*cdp.Node + ctxx := chromedp.FromContext(ctx) + var target string - log.Printf("Caputure Request for %s\n", url) - chromedp.Run(ctx, - chromedp.Navigate(url), - chromedp.Sleep(time.Second*2), - chromedp.CaptureScreenshot(&scrcap), - chromedp.Nodes("a", &nodes, chromedp.ByQueryAll)) + log.Printf("Caputure Request for %s\n", url) + chromedp.Run(ctx, + chromedp.Navigate(url), + chromedp.Sleep(time.Second*2), + chromedp.CaptureScreenshot(&scrcap), + chromedp.Nodes("a", &nodes, chromedp.ByQueryAll)) - fmt.Fprintf(out, "\n\n") + log.Printf("Done with caputure for %s\n", url) } func main() { - ctx, cancel = chromedp.NewContext(context.Background()) - defer cancel() + ctx, cancel = chromedp.NewContext(context.Background()) + defer cancel() + var addr string + flag.StringVar(&addr, "l", ":8080", "Listen address:port, default :8080") + flag.Parse() - http.HandleFunc("/", pageServer) - http.HandleFunc("/wrp.png", imgServer) - log.Printf("Starting http server on :8080\n") - http.ListenAndServe(":8080", nil) + http.HandleFunc("/", pageServer) + http.HandleFunc("/wrp.png", imgServer) + log.Printf("Starting http server on %s\n", addr) + http.ListenAndServe(addr, nil) }