greenhouse/digitalocean_service.go

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	errors "git.sequentialread.com/forest/pkg-errors"
)

type DigitalOceanService struct {
	BaseHTTPService
	APIUrl  string
	APIKey  string
	SSHKeys []ConfigSSHKey
	Region  string
	Image   string
}

type VPSInstance struct {
	ServiceProvider    string
	ProviderInstanceId string
	TenantId           int
	IPV4               string
	IPV6               string
	Bytes              int64
	BytesAllowance     int64
	BytesMonthly       int64
	Created            time.Time
	Deprecated         bool
	Deleted            bool
}

func (i *VPSInstance) GetId() string {
	return fmt.Sprintf("%s-%s", i.ServiceProvider, i.ProviderInstanceId)
}

type DigitalOceanDroplet struct {
	Status   string                      `json:"status"`
	ID       int                         `json:"id"`
	Networks DigitalOceanDropletNetworks `json:"networks"`
	Created  string                      `json:"created_at"`
	SizeSlug string                      `json:"size_slug"`
}

type DigitalOceanDropletNetworks struct {
	IPv4 []DigitalOceanIPv4Network `json:"v4"`
	IPv6 []DigitalOceanIPv6Network `json:"v6"`
}

type DigitalOceanCreateDropletRequest struct {
	Name     string   `json:"name"`
	Region   string   `json:"region"`
	Image    string   `json:"image"`
	Size     string   `json:"size"`
	SSHKeys  []int    `json:"ssh_keys"`
	Backups  bool     `json:"backups"`
	IPV6     bool     `json:"ipv6"`
	UserData string   `json:"user_data"`
	Tags     []string `json:"tags"`
}

type DigitalOceanIPv4Network struct {
	IPAddress string `json:"ip_address"`
	Netmask   string `json:"netmask"`
	Gateway   string `json:"gateway"`
	Type      string `json:"type"`
}

type DigitalOceanIPv6Network struct {
	IPAddress string `json:"ip_address"`
	Netmask   int    `json:"netmask"`
	Gateway   string `json:"gateway"`
	Type      string `json:"type"`
}

type DigitalOceanSSHKey struct {
	Id          int    `json:"id,omitempty"`
	Fingerprint string `json:"fingerprint,omitempty"`
	PublicKey   string `json:"public_key,omitempty"`
	Name        string `json:"name,omitempty"`
}

const DEFAULT_INSTANCE_SIZE = "s-1vcpu-1gb"
const DEFAULT_INSTANCE_MONTHLY_BYTES = TERABYTE

func NewDigitalOceanService(config *Config) *DigitalOceanService {
	toReturn := &DigitalOceanService{
		APIUrl:  "https://api.digitalocean.com",
		APIKey:  config.DigitalOceanAPIKey,
		SSHKeys: config.DigitalOceanSSHAuthorizedKeys,
		Region:  config.DigitalOceanRegion,
		Image:   config.DigitalOceanImage,
	}

	toReturn.ClientFactory = func() (*http.Client, *time.Time, error) {
		return &http.Client{Timeout: 300 * time.Second}, nil, nil
	}

	return toReturn
}

func (service *DigitalOceanService) List() (map[string]*VPSInstance, error) {

	responseBytes, err := service.DigitalOceanHTTP("GET", "/v2/droplets?tag_name=greenhouse-worker", nil)
	if err != nil {
		return nil, err
	}

	var responseObject struct {
		Droplets []DigitalOceanDroplet `json:"droplets"`
	}
	err = json.Unmarshal(responseBytes, &responseObject)
	if err != nil {
		log.Printf("\n\nresponse from /v2/droplets:\n %s\n\n", string(responseBytes))
		return nil, errors.Wrap(err, "JSON parse error when GET-ing /v2/droplets on digitalocean API")
	}

	if responseObject.Droplets == nil {
		log.Printf("\n\nresponse from /v2/droplets:\n %s\n\n", string(responseBytes))
		return nil, errors.New("response does not have droplets property")
	}

	toReturn := map[string]*VPSInstance{}
	for _, droplet := range responseObject.Droplets {
		instance, err := service.DropletToInstance(droplet)
		if err != nil {
			return nil, errors.Wrap(err, "mapping error when GET-ing /v2/droplets on digitalocean API")
		}
		toReturn[instance.GetId()] = instance
	}

	return toReturn, nil
}

func (service *DigitalOceanService) Get(id string) (*VPSInstance, error) {

	responseBytes, err := service.DigitalOceanHTTP("GET", fmt.Sprintf("/v2/droplets/%s", id), nil)
	if err != nil {
		return nil, err
	}

	var responseObject struct {
		Droplet DigitalOceanDroplet `json:"droplet"`
	}
	err = json.Unmarshal(responseBytes, &responseObject)
	if err != nil {
		log.Printf("\n\nresponse from /v2/droplets/:\n %s\n\n", string(responseBytes))
		return nil, errors.Wrap(err, "JSON parse error when GET-ing /v2/droplets on digitalocean API")
	}
	if responseObject.Droplet.ID == 0 {
		log.Printf("\n\nresponse from /v2/droplets:\n %s\n\n", string(responseBytes))
		return nil, errors.New("response does not have droplet.id property")
	}

	return service.DropletToInstance(responseObject.Droplet)
}

func (service *DigitalOceanService) Create(roleName, userData string) (*VPSInstance, error) {

	keyIds, err := service.EnsureDigitalOceanSSHKeyIds()
	if err != nil {
		return nil, err
	}

	randomSuffixBuffer := make([]byte, 4)
	rand.Read(randomSuffixBuffer)
	randomSuffix := hex.EncodeToString(randomSuffixBuffer)

	bodyBytes, err := json.Marshal(DigitalOceanCreateDropletRequest{
		Name:     fmt.Sprintf("greenhouse-%s-%s", roleName, randomSuffix),
		SSHKeys:  keyIds,
		Region:   service.Region,
		Image:    service.Image,
		Size:     DEFAULT_INSTANCE_SIZE,
		Backups:  false,
		IPV6:     true,
		UserData: userData,
		Tags:     []string{"greenhouse-worker"},
	})

	if err != nil {
		return nil, errors.Wrap(err, "JSON serialization error when calling /v2/account/keys on digitalocean API")
	}

	responseBytes, err := service.DigitalOceanHTTP("POST", "/v2/droplets", bytes.NewBuffer(bodyBytes))
	if err != nil {
		return nil, err
	}

	var responseObject struct {
		Droplet DigitalOceanDroplet `json:"droplet"`
	}
	err = json.Unmarshal(responseBytes, &responseObject)
	if err != nil {
		log.Printf("\n\nresponse from POST /v2/droplets:\n %s\n\n", string(responseBytes))
		return nil, errors.Wrap(err, "JSON parse error when POSTing to /v2/droplets on digitalocean API")
	}

	instance, err := service.DropletToInstance(responseObject.Droplet)
	if err != nil {
		return nil, errors.Wrap(err, "mapping error when POSTing to /v2/droplets on digitalocean API")
	}

	return instance, nil
}

func (service *DigitalOceanService) EnsureDigitalOceanSSHKeyIds() ([]int, error) {

	responseBytes, err := service.DigitalOceanHTTP("GET", "/v2/account/keys", nil)
	if err != nil {
		return nil, err
	}
	var responseObject struct {
		SSHKeys []DigitalOceanSSHKey `json:"ssh_keys"`
	}
	err = json.Unmarshal(responseBytes, &responseObject)
	if err != nil {
		return nil, errors.Wrap(err, "JSON parse error when calling /v2/account/keys on digitalocean API")
	}

	digitalOceanIdByKeyContent := map[string]int{}
	for _, sshKey := range responseObject.SSHKeys {
		digitalOceanIdByKeyContent[strings.TrimSpace(sshKey.PublicKey)] = sshKey.Id
	}

	// create all keys that have not been created yet.
	for _, key := range service.SSHKeys {
		if _, has := digitalOceanIdByKeyContent[strings.TrimSpace(key.PublicKey)]; !has {
			keyId, err := service.CreateDigitalOceanSSHKey(key)
			if err != nil {
				return nil, err
			}
			digitalOceanIdByKeyContent[strings.TrimSpace(key.PublicKey)] = keyId
		}
	}

	toReturn := make([]int, len(service.SSHKeys))

	for i, key := range service.SSHKeys {
		toReturn[i] = digitalOceanIdByKeyContent[strings.TrimSpace(key.PublicKey)]
	}

	return toReturn, nil
}

func (service *DigitalOceanService) CreateDigitalOceanSSHKey(publicKey ConfigSSHKey) (int, error) {

	bodyBytes, err := json.Marshal(DigitalOceanSSHKey{
		Name:      publicKey.Name,
		PublicKey: publicKey.PublicKey,
	})
	if err != nil {
		return -1, errors.Wrap(err, "JSON serialization error when calling /v2/account/keys on digitalocean API")
	}

	responseBytes, err := service.DigitalOceanHTTP("POST", "/v2/account/keys", bytes.NewBuffer(bodyBytes))
	if err != nil {
		return -1, err
	}
	var responseObject DigitalOceanSSHKey
	err = json.Unmarshal(responseBytes, &responseObject)
	if err != nil {
		return -1, errors.Wrap(err, "JSON parse error when calling /v2/account/keys on digitalocean API")
	}

	return responseObject.Id, nil
}

func (service *DigitalOceanService) DigitalOceanHTTP(method string, path string, body io.Reader) ([]byte, error) {
	return service.MyHTTP200(
		method,
		fmt.Sprintf("%s%s", service.APIUrl, path),
		body,
		func(request *http.Request) {
			request.Header.Add("Authorization", fmt.Sprintf("Bearer %s", service.APIKey))
			if body != nil {
				request.Header.Add("Content-Type", "application/json")
			}
		},
	)
}

func (service *DigitalOceanService) DropletToInstance(droplet DigitalOceanDroplet) (*VPSInstance, error) {
	ipv4 := ""
	ipv6 := ""
	if droplet.Networks.IPv4 != nil {
		for _, network := range droplet.Networks.IPv4 {
			if network.Type == "public" {
				ipv4 = network.IPAddress
			}
		}
	}
	if droplet.Networks.IPv6 != nil {
		for _, network := range droplet.Networks.IPv6 {
			if network.Type == "public" {
				ipv6 = network.IPAddress
			}
		}
	}
	created, err := time.Parse(time.RFC3339, droplet.Created)
	if err != nil {
		return nil, errors.Wrap(err, "invalid date")
	}
	bytesMonthly := DEFAULT_INSTANCE_MONTHLY_BYTES
	if droplet.SizeSlug != DEFAULT_INSTANCE_SIZE {
		panic(fmt.Errorf("unknown monthly bandwidth for digital ocean instance %d with size_slug %s", droplet.ID, droplet.SizeSlug))
	}

	return &VPSInstance{
		ServiceProvider:    "digitalocean",
		ProviderInstanceId: strconv.Itoa(droplet.ID),
		IPV4:               ipv4,
		IPV6:               ipv6,
		Created:            created,
		BytesMonthly:       bytesMonthly,
	}, nil
}

func (service *DigitalOceanService) GetSSHHostKeysFileScript() string {
	return `
FILE_PATH="greenhouse/known-hosts/digitalocean-$(curl -sS http://169.254.169.254/metadata/v1/id)"
IP_ADDRESS="$(curl -sS http://169.254.169.254/metadata/v1/interfaces/public/0/ipv4/address)"
CONTENT="$( cat /etc/ssh/ssh_host_ecdsa_key.pub /etc/ssh/ssh_host_ed25519_key.pub /etc/ssh/ssh_host_rsa_key.pub | awk "{ print \"$IP_ADDRESS \""'$1" "$2'" }" )"
  `
}