webdriver/driver.go

package webdriver

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os/exec"
	"strconv"
)

type Driver struct {
	server *exec.Cmd

	host      string
	port      uint16
	sessionID string
}

func NewDriver() (*Driver, error) {
	host := "127.0.0.1"
	port := uint16(4444)

	ctx := context.Background()
	server := exec.CommandContext(ctx, "geckodriver", "--host", host, "--port", strconv.Itoa(int(port)))

	stdout, err := server.StdoutPipe()
	if err != nil {
		return nil, err
	}

	if err := server.Start(); err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(stdout)
	scanner.Scan()

	return &Driver{
		server: server,
		host:   host,
		port:   port,
	}, nil
}

func UseDriver(host string, port uint16) *Driver {
	return &Driver{
		host: host,
		port: port,
	}
}

func (d *Driver) Close() error {
	if d.server == nil {
		return nil
	}

	if d.sessionID != "" {
		if err := d.CloseSession(); err != nil {
			return err
		}
	}

	if err := d.server.Cancel(); err != nil {
		return err
	}

	if err := d.server.Wait(); err != nil {
		return err
	}

	return nil
}

func (d *Driver) buildSessionUrl(path string) string {
	uri, err := url.JoinPath(fmt.Sprintf("http://%s:%d/session/%s", d.host, d.port, d.sessionID), path)
	if err != nil {
		panic(err)
	}
	return uri
}

func (d *Driver) buildUrl(path string) string {
	uri, err := url.JoinPath(fmt.Sprintf("http://%s:%d", d.host, d.port), path)
	if err != nil {
		panic(err)
	}
	return uri
}

func (d *Driver) DoRawSessionRequest(method string, path string, body io.Reader) (*http.Response, error) {
	req, err := http.NewRequest(method, d.buildSessionUrl(path), body)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Content-Type", "application/json")

	return http.DefaultClient.Do(req)
}

func (d *Driver) DoRawRequest(method string, path string, body io.Reader) (*http.Response, error) {
	req, err := http.NewRequest(method, d.buildUrl(path), body)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Content-Type", "application/json")

	return http.DefaultClient.Do(req)
}

func (d *Driver) DoSessionRequest(method string, path string, body map[string]interface{}) (map[string]interface{}, error) {
	if body == nil {
		body = map[string]interface{}{}
	}

	data, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}

	resp, err := d.DoRawSessionRequest(method, path, bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	respBody := map[string]interface{}{}
	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
		return nil, err
	}

	return respBody, nil
}

func (d *Driver) DoRequest(method string, path string, body map[string]interface{}) (map[string]interface{}, error) {
	if body == nil {
		body = map[string]interface{}{}
	}

	data, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}

	resp, err := d.DoRawRequest(method, path, bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	respBody := map[string]interface{}{}
	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
		return nil, err
	}

	return respBody, nil
}

func (d *Driver) NewSession() error {
	data, err := d.DoRequest(http.MethodPost, "/session", nil)
	if err != nil {
		return err
	}

	d.sessionID = data["value"].(map[string]interface{})["sessionId"].(string)
	return nil
}

func (d *Driver) CloseSession() error {
	_, err := d.DoSessionRequest(http.MethodDelete, "", nil)
	if err != nil {
		return err
	}

	d.sessionID = ""
	return nil
}