Code to accompany https://pdxjohnny.github.io/digital-ocean-python-scripting/

do.py

# Documentation: https://docs.python.org/3.8/py-modindex.html
import os
import time
import socket
import pathlib
import getpass

# Documentation: https://github.com/koalalorenzo/python-digitalocean
import digitalocean

# Documentation: https://parallel-ssh.readthedocs.io/en/latest/
import pssh.utils
import pssh.clients
import gevent


def digitalocean_ensure_ssh_key_present(
    manager, public_ssh_key_path=pathlib.Path("~", ".ssh", "id_rsa.pub").expanduser()
):
    """
    Make sure that the SSH public key on our system is present within our
    DigitalOcean account

    We require a pathlib.Path object used to reference the file path of the
    public key file on disk. It's in our home directory which we reference
    using the `~` character. When we use that character we MUST call
    the expanduser() method on the Path object to replace that character with
    the correct path to the home directory.
    """
    # We'll be using the private key to connect to the server. It's the same
    # file path as the public key only without the .pub suffix
    private_ssh_key_path = public_ssh_key_path.with_suffix("")
    # The key's name will be our username plus the name of the machine we're on
    key_name = getpass.getuser() + "_" + socket.gethostname()
    # Check if the key exists already
    for ssh_key in manager.get_all_sshkeys():
        if ssh_key.name == key_name:
            return private_ssh_key_path, ssh_key
    # We then read the contents of the file into a variable
    public_ssh_key_contents = public_ssh_key_path.read_text()
    # Create a digitalocean.SSHKey object. Reuse the token from the manager
    public_ssh_key = digitalocean.SSHKey(
        token=manager.token, name=key_name, public_key=public_ssh_key_contents,
    )
    # Call the create() method on the key object to create the key
    public_ssh_key.create()
    # Return the key
    return private_ssh_key_path, public_ssh_key


def print_image_slugs(manager):
    """
    Call this function to have a list of image descriptions and their slug's
    printed to the terminal.
    """
    for image in manager.get_global_images():
        print(f"{image.slug}: {image.description}")


def print_region_slugs(manager):
    """
    Call this function to have a list of data centers and their slug's printed
    to the terminal.
    """
    for region in manager.get_all_regions():
        print(f"{region.slug}: {region.name}")


def print_size_slugs(manager):
    """
    Call this function to have a list of VM sizes and their slug's printed
    to the terminal.
    """
    for size in manager.get_all_sizes():
        print(f"{size.slug:<15}", end=" ")
        # Print out the following attributes from the size object
        for attr in [
            "vcpus",
            "memory",
            "disk",
            "transfer",
            "price_monthly",
            "price_hourly",
            "regions",
        ]:
            value = getattr(size, attr)
            print(f"{attr}: {value!s:<8}", end=" ")
        print()


def main():
    # Get the API access token for our account from the environment variables
    if not "DIGITALOCEAN_ACCESS_TOKEN" in os.environ:
        raise Exception(
            "DIGITALOCEAN_ACCESS_TOKEN must be set as an environment variable"
        )
    digitalocean_access_token = os.environ["DIGITALOCEAN_ACCESS_TOKEN"]
    # Create an instance of the manager object to interact with the Digital
    # Ocean API.
    manager = digitalocean.Manager(token=digitalocean_access_token)

    # A "slug" is the unique string we use to refer to a particular VM iamge
    # Use print_image_slugs(manager) to see more options
    image_slug = "docker-20-04"
    # The region is the data center we want to run the VM in
    # Use print_region_slugs(manager) to see more options
    region_slug = "sfo2"
    # The VM size is how big of a machine we are asking for
    # Use print_size_slugs(manager) to see more options
    size_slug = "s-1vcpu-1gb"

    # The name of this project
    project_name = "my-test-app"

    # The template to use for the VMs names. Use f string formating to add the
    # variables we already have defined into the name
    name_template = f"{project_name}-{region_slug}-{size_slug}"
    # The str.format(*args) method will look for any {} and replace their
    # contents with whatever ever args were passed. We need to append it after
    # we have done any f string formating so that Python doesn't try to fill the
    # empty {} at time of f string formatting. We only want it filled later when
    # we call str.format()
    name_template += "-{}"

    # The number of VMs to create
    num_vms = 2

    # Make sure we have our machine's ssh key registered with Digital Ocean
    private_ssh_key_path, public_ssh_key = digitalocean_ensure_ssh_key_present(manager)

    # Create a mapping of droplet names to their objects for all VMs we have
    all_vms = {vm.name: vm for vm in manager.get_all_droplets()}

    # Create a tag name using the project name to associate the droplets with
    # this project. Spaces are not allowed
    project_tag_name = f"project:{project_name}"
    # Create a digitalocean.Tag object. Reuse the token from the manager
    project_tag = digitalocean.Tag(token=manager.token, name=project_tag_name)
    # Call the create() method on the tag object to create the tag
    project_tag.create()

    # Create a VM if it doesn't already exist
    for i in range(0, num_vms):
        # The name of this VM
        name = name_template.format(i)
        # Skip creation if it exists
        if name in all_vms:
            continue
        # Create a digitalocean.Droplet object. Reuse the token from the manager
        vm = digitalocean.Droplet(
            token=manager.token,
            name=name,
            region=region_slug,
            image=image_slug,
            size_slug=size_slug,
            ssh_keys=[public_ssh_key],
            backups=False,
            tags=[project_tag_name],
        )
        # Call the create() method on the droplet object to create the droplet
        vm.create()

    # Create a mapping of all the VMs associated with this project. Do this
    # until the mapping contains the same number of VMs that should exist
    project_vms = {
        vm.name: vm for vm in manager.get_all_droplets(tag_name=project_tag_name)
    }
    print(project_vms)
    while len(project_vms) != num_vms:
        time.sleep(1)
        project_vms = {
            vm.name: vm for vm in manager.get_all_droplets(tag_name=project_tag_name)
        }
        print(project_vms)

    # Creating VMs is finished. Now we'll ssh into the project's VMs to setup
    # the software we want running on them

    # The user we'll be loging in as. Different images might do this differently
    ssh_user = "root"
    # Create a list of IP addresses we'll be ssh'ing into
    ssh_hosts = [vm.ip_address for vm in project_vms.values()]
    # The SSH client object doesn't know what to do with a pathlib.Path object
    # if it's given one as the private key (pkey). Therefore, we need to resolve
    # the private key pathlib.Path object to find it's absolute path, the path
    # from the root directory to the file. We then need to convert it from a
    # pathlib.Path object to a string by calling the str function
    private_ssh_key_path_as_string = str(private_ssh_key_path.resolve())
    # SSH private keys are typically proctected on disk by encrypting them with
    # a password. We ask the user for their password for this key here so that
    # we can unlock the key and use it to log in to the servers.
    private_ssh_key_password = getpass.getpass(
        prompt=f"SSH Private Key ({private_ssh_key_path_as_string}) Password: "
    )
    # Create an SSH client which we'll use to access all the VMs in parallel.
    # Specify that we want to use our private key to connect.
    client = pssh.clients.ParallelSSHClient(
        ssh_hosts,
        user=ssh_user,
        pkey=private_ssh_key_path_as_string,
        password=private_ssh_key_password,
    )
    # Print output of commands to terminal as they run
    pssh.utils.enable_host_logger()

    # Define a list of commands we want to run on all the hosts in parallel
    cmds = [
        # Stop all running containers
        "docker kill $(docker ps -qa)",
        # Delete all containers
        "docker rm $(docker ps -qa)",
    ]
    # Run each command on all hosts
    for cmd in cmds:
        client.run_command(cmd)
        client.join(consume_output=True)

    # Create a file
    index_html_path = pathlib.Path("index.html")
    # Put something in the file
    index_html_path.write_text("Hello World")
    # Copy files over to all machines
    cmds = client.copy_file(str(index_html_path), "html_directory/index.html")
    gevent.joinall(cmds, raise_error=True)

    # Define more commands we want to run on all the hosts in parallel
    cmds = [
        # Start a new background container running our web server
        "docker run -d -p 80:8080 -v $PWD/html_directory:/var/www/html -w /var/www/html python:3.8 python -m http.server 8080",
    ]
    # Run each command on all hosts
    for cmd in cmds:
        client.run_command(cmd)
        client.join(consume_output=True)


if __name__ == "__main__":
    main()