do_floatingip_update.py

#!/usr/bin/env python3

import json
import logging
import os
import requests
import sys
import time

base_url = "https://api.digitalocean.com/v2"


def get_droplet_id(headers, tag_name, timeout=5):
    try:
        response = requests.get(
            f'{base_url}/droplets?tag_name={tag_name}',
            headers=headers,
            timeout=timeout,
        )

        status = response.status_code
        if status == 200:
            resp_json = response.json()
            droplet_id = resp_json['droplets'][0]['id']
            return droplet_id
        else:
            logging.error("Unable to get droplet ID. Status code: " +
                          str(status))
            sys.exit(1)

    except Exception as e:
        logging.error("Unknown error while getting Algo droplet ID: ", e)
        sys.exit(1)


def get_algo_floating_ip(headers, droplet_id, timeout=5):
    try:
        response = requests.get(
            f'{base_url}/floating_ips',
            headers=headers,
            timeout=timeout,
        )

        status = response.status_code
        if response.status_code == 200:
            resp_json = response.json()

            for ip in resp_json['floating_ips']:
                # See if a given floating IP has a droplet ID, if not, continue to the next one
                try:
                    floating_ip_droplet_id = ip['droplet']['id']
                    if floating_ip_droplet_id == droplet_id:
                        return ip['ip']
                except TypeError:
                    continue
        else:
            logging.error("Error getting Algo floating IP. Status code: " +
                          str(status))
            sys.exit(1)

    except Exception as e:
        logging.error("Unknown error while getting Algo floating IP: ", e)
        sys.exit(1)


def delete_algo_floating_ip(headers, droplet_id, timeout=5):
    try:
        ip = get_algo_floating_ip(headers, droplet_id)
        response = requests.delete(
            f'{base_url}/floating_ips/{ip}',
            headers=headers,
            timeout=timeout,
        )

        status = response.status_code
        if status == 204:
            logging.info("Deleted Algo floating IP.")
        else:
            logging.error("Error deleting Algo floating IP. Status code: " +
                          str(status))

    except TypeError:
        logging.error("No Algo floating IP to delete.")
    except Exception as e:
        logging.error("Unknown error while deleting Algo IP: ", e)
        sys.exit(1)


def create_regional_floating_ip(headers, timeout=5):
    try:
        payload = {"region": "nyc3"}
        response = requests.post(
            f'{base_url}/floating_ips',
            json=payload,
            headers=headers,
            timeout=timeout,
        )

        status = response.status_code
        if status == 202:
            resp_json = response.json()
            ip = resp_json['floating_ip']['ip']
            logging.info("Created regional floating IP.")
            logging.debug(f"IP: {ip}")
            return ip
        else:
            logging.error(
                "Error creating regional floating IP, are you at the limit of 3? Status code: "
                + str(status))
            sys.exit(1)

    except Exception as e:
        logging.error(
            "Unknown error while creating Algo regional floating IP: ", e)
        sys.exit(1)


def assign_algo_floating_ip(headers, ip, droplet_id, timeout=5):
    try:
        payload = {"type": "assign", "droplet_id": droplet_id}
        response = requests.post(
            f'{base_url}/floating_ips/{ip}/actions',
            json=payload,
            headers=headers,
            timeout=timeout,
        )

        status = response.status_code
        if status == 201:
            logging.info("Assigned floating IP to Algo.")
        else:
            logging.error(
                "Error assigning a floating IP to Algo. Status code: " +
                str(status))
            logging.error(response.json())
            sys.exit(1)

    except Exception as e:
        logging.error("Unknown error while assigning Algo floating IP: ", e)
        sys.exit(1)


def run():
    logging.basicConfig(
        format='%(asctime)s [%(levelname)s]: %(message)s',
        datefmt='%m/%d/%Y %I:%M:%S %p',
        level=logging.DEBUG)
    if "DO_API_TOKEN" in os.environ:
        logging.info("Running Algo floating IP update.")

        pat = os.environ['DO_API_TOKEN']
        tag_name = "Environment:Algo"
        headers = {
            "Authorization": "Bearer " + pat,
        }
        droplet_id = get_droplet_id(headers, tag_name)
        new_ip = create_regional_floating_ip(headers)
        delete_algo_floating_ip(headers, droplet_id)
        time.sleep(3)
        assign_algo_floating_ip(headers, new_ip, droplet_id)

        logging.info("Completed Algo floating IP update.")
    else:
        logging.error(
            "Unable to get Digital Ocean API token from the DO_API_TOKEN environment variable."
        )
        sys.exit(1)


if __name__ == '__main__':
    run()