From 35fb68da38e81aec597608968ff6aa4d4e02b49e Mon Sep 17 00:00:00 2001
From: Phillip Alexander <pal@jobindex.dk>
Date: Thu, 25 Jan 2024 07:17:06 +0100
Subject: [PATCH] Add action to return page as a PDF

---
 browser.go        | 14 ++++++++++++++
 cmd/decap/main.go | 37 +++++++++++++++++++++++++------------
 query.go          | 28 ++++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/browser.go b/browser.go
index 73259e5..59a191e 100644
--- a/browser.go
+++ b/browser.go
@@ -282,6 +282,20 @@ func outerHTML(out *[]string) chromedp.ActionFunc {
 	}
 }
 
+func printToPDF(buf *[]byte, margins []float64) chromedp.ActionFunc {
+	return func(ctx context.Context) error {
+		if len(margins) != 4 {
+			return fmt.Errorf("expected four margins (top, right, bottom, left)")
+		}
+		var err error
+		p := page.PrintToPDF()
+		p = p.WithMarginTop(margins[0]).WithMarginBottom(margins[2])
+		p = p.WithMarginLeft(margins[3]).WithMarginRight(margins[1])
+		*buf, _, err = p.Do(ctx)
+		return err
+	}
+}
+
 func removeElements(sel string) chromedp.ActionFunc {
 	return func(ctx context.Context) error {
 		cmd := fmt.Sprintf("document.querySelectorAll('%s').forEach(e => e.remove());", sel)
diff --git a/cmd/decap/main.go b/cmd/decap/main.go
index cbea3bd..2a49175 100644
--- a/cmd/decap/main.go
+++ b/cmd/decap/main.go
@@ -105,13 +105,13 @@ func browseHandler(w http.ResponseWriter, req *http.Request) {
 
 	// execute query
 
+	err_status := http.StatusInternalServerError
 	var res *decap.Result
 	res, err = dec.Execute()
 	if err != nil {
 		// TODO: Propagate HTTP status properly
-		status := http.StatusInternalServerError
-		msg := fmt.Sprintf("%s: %s", http.StatusText(status), err)
-		http.Error(w, msg, status)
+		msg := fmt.Sprintf("%s: %s", http.StatusText(err_status), err)
+		http.Error(w, msg, err_status)
 		return
 	}
 
@@ -121,21 +121,34 @@ func browseHandler(w http.ResponseWriter, req *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 		err = json.NewEncoder(w).Encode(res)
 		if err != nil {
-			status := http.StatusInternalServerError
-			msg := fmt.Sprintf("%s: %s", http.StatusText(status), "Couldn't encode response")
-			http.Error(w, msg, status)
-			return
+			msg := fmt.Sprintf("%s: %s", http.StatusText(err_status), "Couldn't encode response")
+			http.Error(w, msg, err_status)
 		}
+		return
+	case "pdf":
+		w.Header().Set("Content-Type", "application/pdf")
+		_, err = w.Write(res.PDFBuffer())
+		if err != nil {
+			msg := fmt.Sprintf("%s: %s",
+				http.StatusText(err_status), "Couldn't write response bytes")
+			http.Error(w, msg, err_status)
+		}
+		return
 	case "png":
 		w.Header().Set("Content-Type", "image/png")
-		_, err := w.Write(res.ImgBuffer())
+		_, err = w.Write(res.ImgBuffer())
 		if err != nil {
-			status := http.StatusInternalServerError
 			msg := fmt.Sprintf("%s: %s",
-				http.StatusText(status), "Couldn't write response bytes")
-			http.Error(w, msg, status)
-			return
+				http.StatusText(err_status), "Couldn't write response bytes")
+			http.Error(w, msg, err_status)
 		}
+		return
+	default:
+		fmt.Fprintf(os.Stderr, `unknown results type "%s"`, res.Type())
+		msg := fmt.Sprintf(`%s: Unknown result type "%s"`,
+			http.StatusText(err_status), res.Type())
+		http.Error(w, msg, err_status)
+		return
 	}
 }
 
diff --git a/query.go b/query.go
index e199c6e..838dd06 100644
--- a/query.go
+++ b/query.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/url"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -33,19 +34,28 @@ type Result struct {
 	TabID    string     `json:"tab_id"`
 	WindowID string     `json:"window_id"`
 	img      []byte
+	pdf      []byte
 }
 
 func (res *Result) Type() string {
-	if len(res.img) != 0 {
+	switch {
+	case len(res.pdf) != 0:
+		return "pdf"
+	case len(res.img) != 0:
 		return "png"
+	default:
+		return "json"
 	}
-	return "json"
 }
 
 func (res *Result) ImgBuffer() []byte {
 	return res.img
 }
 
+func (res *Result) PDFBuffer() []byte {
+	return res.pdf
+}
+
 type QueryBlock struct {
 	Actions    []Action `json:"actions"`
 	Repeat     *int     `json:"repeat"`
@@ -416,6 +426,20 @@ func (r *Request) parseAction(xa Action) error {
 		}
 		r.appendActions(outerHTML(&r.res.Out[r.pos]))
 
+	case "print_to_pdf":
+		margins := make([]float64, 4)
+		if err = xa.MustArgCount(0, 4); err != nil {
+			return err
+		}
+		for i, v := range xa.Args() {
+			if margins[i], err = strconv.ParseFloat(v, 64); err != nil {
+				msg := "print_to_pdf: expected floating point margins"
+				return fmt.Errorf("%s: %w", msg, err)
+			}
+		}
+
+		r.appendActions(printToPDF(&r.res.pdf, margins))
+
 	case "remove":
 		if len(xa.Args()) == 0 {
 			return fmt.Errorf("remove: expected at least one argument")