vkiosk/collect.go

// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"v.io/jiri"
	"v.io/jiri/runutil"
	"v.io/x/lib/cmdline"
)

var (
	displayFlag            string
	resolutionFlag         string
	screenshotIntervalFlag string
	screenshotNameFlag     string
	urlFlag                string
)

func init() {
	cmdCollect.Flags.StringVar(&displayFlag, "display", ":1365", "The value of DISPLAY environment variable for Xvfb.")
	cmdCollect.Flags.StringVar(&resolutionFlag, "resolution", "1920x1080x24", "The resolution string for Xvfb.")
	cmdCollect.Flags.StringVar(&screenshotIntervalFlag, "interval", "5s", "The interval between screenshots.")
	cmdCollect.Flags.StringVar(&screenshotNameFlag, "name", "", "The name of the screenshot file.")
	cmdCollect.Flags.StringVar(&urlFlag, "url", "", "The url to take screenshots for.")
}

var cmdCollect = &cmdline.Command{
	Name:  "collect",
	Short: "Takes screenshots of a given URL in Chrome and stores them in the given export dir",
	Long: `
The collect commands takes screenshots of a given URL in Chrome and stores them
in the given export dir.

To use this command, the following programs need to be installed:
Google Chrome,Xvfb, and Fluxbox.
`,
	Runner: cmdline.RunnerFunc(runCollect),
}

func runCollect(env *cmdline.Env, args []string) error {
	jirix, err := jiri.NewX(env)
	if err != nil {
		return err
	}
	if err := checkPreRequisites(); err != nil {
		return err
	}

	// A tmp dir to store screenshots.
	tmpDir, err := jirix.NewSeq().TempDir("", "vkiosk")
	if err != nil {
		return err
	}
	fmt.Fprintf(jirix.Stdout(), "Tmp screenshot dir: %s\n", tmpDir)

	fmt.Fprintf(jirix.Stdout(), "Starting Xvfb at DISPLAY=%s with resolution %s...\n", displayFlag, resolutionFlag)
	p, err := startXvfb(jirix)
	if err != nil {
		return err
	}

	// Set up cleanup function to remove tmp screenshot dir and kill Xvfb.
	// When Xvfb is killed, the Chrome running in it will also be killed
	// automatically.
	cleanupFn := func() {
		// Ignore all errors.
		jirix.NewSeq().RemoveAll(tmpDir)
		p.Kill()
	}
	defer cleanupFn()

	// Trap SIGTERM and SIGINT signal.
	go func() {
		sigchan := make(chan os.Signal, 1)
		signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)
		<-sigchan
		cleanupFn()
		os.Exit(0)
	}()

	fmt.Fprintf(jirix.Stdout(), "Starting Chrome for %q in Xvfb...\n", urlFlag)
	if err := startChrome(jirix); err != nil {
		return err
	}

	fmt.Fprintf(jirix.Stdout(), "Starting Fluxbox in Xvfb...\n")
	if err := startFluxbox(jirix); err != nil {
		return err
	}

	if err := takeScreenshots(jirix, tmpDir); err != nil {
		return err
	}
	return nil
}

func checkPreRequisites() error {
	programs := []string{"Xvfb", "fluxbox", "google-chrome", "scrot"}
	notInstalled := []string{}
	for _, p := range programs {
		if _, err := exec.LookPath(p); err != nil {
			notInstalled = append(notInstalled, p)
		}
	}
	if len(notInstalled) != 0 {
		return fmt.Errorf("%q not found in PATH", strings.Join(notInstalled, ","))
	}
	return nil
}

// startXvfb starts Xvfb at the given display and screen resolution.
// It provides an virtual X11 environment for Chrome to run in.
func startXvfb(jirix *jiri.X) (*runutil.Handle, error) {
	args := []string{
		displayFlag,
		"-screen",
		"0",
		resolutionFlag,
		"-ac",
	}
	cmd, err := jirix.NewSeq().Start("Xvfb", args...)
	if err != nil {
		return nil, err
	}
	// Wait a little bit to make sure it finishes initialization.
	time.Sleep(time.Second * 5)
	return cmd, nil
}

// startFluxbox starts a light weight windows manager Fluxbox.
// Without it, Chrome can't be maximized or run in kiosk mode properly.
func startFluxbox(jirix *jiri.X) error {
	args := []string{
		"-display",
		displayFlag,
	}
	if _, err := jirix.NewSeq().Start("fluxbox", args...); err != nil {
		return err
	}
	return nil
}

// startChrome starts Chrome in a kiosk mode in the given X11 environment.
func startChrome(jirix *jiri.X) error {
	args := []string{
		"--kiosk",
		urlFlag,
	}
	env := map[string]string{"DISPLAY": displayFlag}
	if _, err := jirix.NewSeq().Env(env).Start("google-chrome", args...); err != nil {
		return err
	}
	return nil
}

// takeScreenshot takes screenshots periodically, and stores them in the given
// export dir.
func takeScreenshots(jirix *jiri.X, tmpDir string) error {
	d, err := time.ParseDuration(screenshotIntervalFlag)
	if err != nil {
		return fmt.Errorf("ParseDuration(%s) failed: %v", screenshotIntervalFlag, err)
	}
	screenshotFile := filepath.Join(tmpDir, screenshotNameFlag)
	scrotArgs := []string{
		screenshotFile,
	}
	s := jirix.NewSeq()
	env := map[string]string{"DISPLAY": displayFlag}
	gsutilArgs := []string{
		"-q",
		"cp",
		screenshotFile,
		exportDirFlag + "/" + screenshotNameFlag,
	}
	ticker := time.NewTicker(d)
	for range ticker.C {
		// Use "scrot" command to take screenshots.
		fmt.Fprintf(jirix.Stdout(), "[%s]: take screenshot to %q...\n", nowTimestamp(), screenshotFile)
		if err := s.Env(env).Last("scrot", scrotArgs...); err != nil {
			fmt.Fprintf(jirix.Stderr(), "%v\n", err)
			continue
		}

		// Store the screenshots to export dir.
		fmt.Fprintf(jirix.Stdout(), "[%s]: copying screenshot to %s...\n", nowTimestamp(), exportDirFlag)
		if strings.HasPrefix(exportDirFlag, "gs://") {
			if err := s.Last("gsutil", gsutilArgs...); err != nil {
				fmt.Fprintf(jirix.Stderr(), "%v\n", err)
			}
		} else {
			if err := s.Rename(screenshotFile, filepath.Join(exportDirFlag, screenshotNameFlag)).Done(); err != nil {
				fmt.Fprintf(jirix.Stderr(), "%v\n", err)
			}
		}
	}
	return nil
}

func nowTimestamp() string {
	return time.Now().Format("20060102 15:04:05")
}