Capture a website screenshot with Go.

// screenshot.go: request a screenshot and save it. Run: go run screenshot.go
package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

// errorEnvelope is what every failed call returns instead of an image.
type errorEnvelope struct {
	Error struct {
		Code      string `json:"code"`
		Message   string `json:"message"`
		RequestID string `json:"requestId"`
	} `json:"error"`
}

func main() {
	query := url.Values{}
	query.Set("url", "https://en.wikipedia.org/wiki/Cartography")
	query.Set("format", "webp")
	query.Set("quality", "80")
	query.Set("viewport_width", "1280")
	query.Set("viewport_height", "800")
	query.Set("device_scale_factor", "2")

	request, err := http.NewRequest(http.MethodGet, "https://api.urlshot.io/v1/screenshot?"+query.Encode(), nil)
	if err != nil {
		log.Fatal(err)
	}
	request.Header.Set("Authorization", "Bearer "+os.Getenv("URLSHOT_API_KEY"))

	client := &http.Client{Timeout: 60 * time.Second}
	response, err := client.Do(request)
	if err != nil {
		log.Fatal(err)
	}
	defer response.Body.Close()

	body, err := io.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}

	// Anything that is not an image is the JSON error envelope.
	if response.StatusCode != http.StatusOK {
		var envelope errorEnvelope
		if err := json.Unmarshal(body, &envelope); err != nil {
			log.Fatal(err)
		}
		log.Fatalf("%s: %s (request %s)", envelope.Error.Code, envelope.Error.Message, envelope.Error.RequestID)
	}

	if err := os.WriteFile("screenshot.webp", body, 0o644); err != nil {
		log.Fatal(err)
	}
	fmt.Println("Wrote screenshot.webp")
}

// After the request: what it cost, and how to refer to it. Header.Get
// matches a name whatever its case.
credits := response.Header.Get("X-Urlshot-Credits-Used") // "1", or "0" from cache
requestID := response.Header.Get("X-Request-ID")
fmt.Printf("request %s used %s credit(s)\n", requestID, credits)

// retry.go: request a screenshot, retrying only when the API says when to.
package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"
)

// errorEnvelope is what every failed call returns instead of an image.
type errorEnvelope struct {
	Error struct {
		Code      string `json:"code"`
		Message   string `json:"message"`
		RequestID string `json:"requestId"`
	} `json:"error"`
}

// screenshot requests an image, retrying only when the API says how long to wait.
func screenshot(query url.Values, attempts int) ([]byte, error) {
	client := &http.Client{Timeout: 60 * time.Second}
	for attempt := 1; ; attempt++ {
		request, err := http.NewRequest(http.MethodGet, "https://api.urlshot.io/v1/screenshot?"+query.Encode(), nil)
		if err != nil {
			return nil, err
		}
		request.Header.Set("Authorization", "Bearer "+os.Getenv("URLSHOT_API_KEY"))

		response, err := client.Do(request)
		if err != nil {
			return nil, err
		}
		body, err := io.ReadAll(response.Body)
		response.Body.Close()
		if err != nil {
			return nil, err
		}
		if response.StatusCode == http.StatusOK {
			return body, nil
		}

		// Only a 429 or 503 carries Retry-After, in seconds. Anything else fails
		// the same way however often it is sent.
		seconds, err := strconv.Atoi(response.Header.Get("Retry-After"))
		if err != nil || attempt == attempts {
			var envelope errorEnvelope
			if err := json.Unmarshal(body, &envelope); err != nil {
				return nil, err
			}
			return nil, fmt.Errorf("%s: %s (request %s)", envelope.Error.Code, envelope.Error.Message, envelope.Error.RequestID)
		}
		time.Sleep(time.Duration(seconds) * time.Second)
	}
}

func main() {
	query := url.Values{}
	query.Set("url", "https://en.wikipedia.org/wiki/Cartography")
	query.Set("format", "webp")

	image, err := screenshot(query, 4)
	if err != nil {
		log.Fatal(err)
	}
	if err := os.WriteFile("screenshot.webp", image, 0o644); err != nil {
		log.Fatal(err)
	}
	fmt.Println("Wrote screenshot.webp")
}

// options.go: request a screenshot and save it. Run: go run options.go
package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

// errorEnvelope is what every failed call returns instead of an image.
type errorEnvelope struct {
	Error struct {
		Code      string `json:"code"`
		Message   string `json:"message"`
		RequestID string `json:"requestId"`
	} `json:"error"`
}

func main() {
	query := url.Values{}
	// Add, not Set, for a list: Set replaces the value, and a list is the
	// parameter repeated, once for each selector.
	query.Set("url", "https://en.wikipedia.org/wiki/Atlas")
	query.Set("full_page", "true")
	query.Set("format", "png")
	query.Add("hide_selectors", ".vector-header-container")
	query.Add("hide_selectors", ".mw-footer")

	request, err := http.NewRequest(http.MethodGet, "https://api.urlshot.io/v1/screenshot?"+query.Encode(), nil)
	if err != nil {
		log.Fatal(err)
	}
	request.Header.Set("Authorization", "Bearer "+os.Getenv("URLSHOT_API_KEY"))

	client := &http.Client{Timeout: 60 * time.Second}
	response, err := client.Do(request)
	if err != nil {
		log.Fatal(err)
	}
	defer response.Body.Close()

	body, err := io.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}

	// Anything that is not an image is the JSON error envelope.
	if response.StatusCode != http.StatusOK {
		var envelope errorEnvelope
		if err := json.Unmarshal(body, &envelope); err != nil {
			log.Fatal(err)
		}
		log.Fatalf("%s: %s (request %s)", envelope.Error.Code, envelope.Error.Message, envelope.Error.RequestID)
	}

	if err := os.WriteFile("atlas.png", body, 0o644); err != nil {
		log.Fatal(err)
	}
	fmt.Println("Wrote atlas.png")
}

// post.go: send the options as a JSON body. Run: go run post.go
package main

import (
	"bytes"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"time"
)

// errorEnvelope is what every failed call returns instead of an image.
type errorEnvelope struct {
	Error struct {
		Code      string `json:"code"`
		Message   string `json:"message"`
		RequestID string `json:"requestId"`
	} `json:"error"`
}

func main() {
	// JSON's own types: true rather than "true", and a list as an array.
	payload, err := json.Marshal(map[string]any{
		"url": "https://en.wikipedia.org/wiki/Atlas",
		"full_page": true,
		"format": "png",
		"hide_selectors": []string{".vector-header-container", ".mw-footer"},
		"custom_css": "#siteNotice { display: none !important }",
	})
	if err != nil {
		log.Fatal(err)
	}

	request, err := http.NewRequest(http.MethodPost, "https://api.urlshot.io/v1/screenshot", bytes.NewReader(payload))
	if err != nil {
		log.Fatal(err)
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer "+os.Getenv("URLSHOT_API_KEY"))

	client := &http.Client{Timeout: 60 * time.Second}
	response, err := client.Do(request)
	if err != nil {
		log.Fatal(err)
	}
	defer response.Body.Close()

	body, err := io.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}

	// Anything that is not an image is the JSON error envelope.
	if response.StatusCode != http.StatusOK {
		var envelope errorEnvelope
		if err := json.Unmarshal(body, &envelope); err != nil {
			log.Fatal(err)
		}
		log.Fatalf("%s: %s (request %s)", envelope.Error.Code, envelope.Error.Message, envelope.Error.RequestID)
	}

	if err := os.WriteFile("atlas.png", body, 0o644); err != nil {
		log.Fatal(err)
	}
}

// server.go: serve previews of chosen pages. Run: go run server.go
package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

// errorEnvelope is what every failed call returns instead of an image.
type errorEnvelope struct {
	Error struct {
		Code      string `json:"code"`
		Message   string `json:"message"`
		RequestID string `json:"requestId"`
	} `json:"error"`
}

// The pages this app shows previews of. Visitors pick a name, never an
// address: a route that took any URL would render any page on your credits.
var pages = map[string]string{
	"home": "https://example.com",
	"atlas": "https://en.wikipedia.org/wiki/Atlas",
}

var client = &http.Client{Timeout: 60 * time.Second}

func preview(w http.ResponseWriter, r *http.Request) {
	target, ok := pages[r.PathValue("page")]
	if !ok {
		http.NotFound(w, r)
		return
	}

	query := url.Values{}
	query.Set("url", target)
	query.Set("viewport_width", "1280")
	query.Set("format", "webp")
	query.Set("cache_ttl", "3600")

	request, err := http.NewRequestWithContext(r.Context(), http.MethodGet, "https://api.urlshot.io/v1/screenshot?"+query.Encode(), nil)
	if err != nil {
		http.Error(w, "preview unavailable", http.StatusInternalServerError)
		return
	}
	request.Header.Set("Authorization", "Bearer "+os.Getenv("URLSHOT_API_KEY"))

	response, err := client.Do(request)
	if err != nil {
		log.Print(err)
		http.Error(w, "preview unavailable", http.StatusBadGateway)
		return
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		var envelope errorEnvelope
		if err := json.NewDecoder(response.Body).Decode(&envelope); err == nil {
			log.Printf("%s: %s (request %s)", envelope.Error.Code, envelope.Error.Message, envelope.Error.RequestID)
		}
		http.Error(w, "preview unavailable", http.StatusBadGateway)
		return
	}

	// Streamed straight through: the image is never held in memory whole.
	w.Header().Set("Content-Type", response.Header.Get("Content-Type"))
	w.Header().Set("Cache-Control", "public, max-age=3600")
	if _, err := io.Copy(w, response.Body); err != nil {
		log.Print(err)
	}
}

func main() {
	http.HandleFunc("GET /previews/{page}", preview)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/url"
	"os"
)

func signedURL(options map[string]string) string {
	query := url.Values{}
	query.Set("key", os.Getenv("URLSHOT_KEY"))
	for name, value := range options {
		query.Set(name, value)
	}
	// Encode sorts by key. That is fine: what matters is signing exactly the string you send.
	encoded := query.Encode()

	mac := hmac.New(sha256.New, []byte(os.Getenv("URLSHOT_SIGNING_SECRET")))
	mac.Write([]byte(encoded))
	signature := hex.EncodeToString(mac.Sum(nil))

	return "https://api.urlshot.io/v1/screenshot?" + encoded + "&signature=" + signature
}

func main() {
	fmt.Println(signedURL(map[string]string{"url": "https://example.com", "viewport_width": "1200"}))
}