digital_ocean.py

import digitalocean
import random
import asyncio
import functools
from datetime import datetime


class DropletManager:
    def __init__(self, thread_pool):
        print("creating DropletManager")
        self.thread_pool = thread_pool
        self.loop = asyncio.get_event_loop()
        self.token = "TOKEN HERE"
        self.manager = digitalocean.Manager(token=self.token)
        self.droplets = None

    async def initialize_manager(self):

        self.droplets = await self.get_all_droplet_ips()

        self.min_droplets = 6

        await self.ensure_min_droplets()

    async def get_all_droplet_ips(self):
        print("getting all droplets initially")
        returnable = {}
        my_droplets = await self.loop.run_in_executor(
            self.thread_pool,
            functools.partial(self.manager.get_all_droplets, tag_name="proxy"),
        )
        for droplet in my_droplets:
            await self.loop.run_in_executor(
                self.thread_pool, functools.partial(droplet.load)
            )

            returnable[droplet.id] = {
                "ip_address": droplet.ip_address,
                "status": droplet.status,
                "region": droplet.region.get("slug"),
            }
        return returnable

    async def delete_droplet(self, droplet_id):
        print(f"deleting droplet by ID {droplet_id}")
        my_droplet = await self.loop.run_in_executor(
            self.thread_pool, functools.partial(self.manager.get_droplet, droplet_id)
        )
        await self.loop.run_in_executor(
            self.thread_pool, functools.partial(my_droplet.destroy)
        )
        del self.droplets[droplet_id]

        await self.ensure_min_droplets()

    async def create_droplet(self):
        print(f"creating droplet...")
        keys = await self.loop.run_in_executor(
            self.thread_pool, functools.partial(self.manager.get_all_sshkeys)
        )

        region, info_blob = await self.get_droplet_create_info()

        droplet = digitalocean.Droplet(
            token=self.token,
            name=f"{info_blob.get('name')}-{int(datetime.timestamp(datetime.utcnow()))}",
            tags=["proxy"],
            region=region,
            image=info_blob.get("id"),
            ssh_keys=keys,
            size_slug="s-1vcpu-1gb",
        )

        await self.loop.run_in_executor(
            self.thread_pool, functools.partial(droplet.create)
        )

        await asyncio.sleep(5)

        await self.loop.run_in_executor(
            self.thread_pool, functools.partial(droplet.load)
        )

        while droplet.status != "active":
            print(
                f"droplet w/ IP {droplet.ip_address} on region {droplet.region.get('slug')} not active yet, waiting..."
            )
            await asyncio.sleep(5)
            await self.loop.run_in_executor(
                self.thread_pool, functools.partial(droplet.load)
            )

        self.droplets[droplet.id] = {
            "ip_address": droplet.ip_address,
            "status": droplet.status,
            "region": droplet.region.get("slug"),
        }

        print(
            f"droplet created - IP:{self.droplets[droplet.id].get('ip_address')}::{self.droplets[droplet.id].get('status')}::{self.droplets[droplet.id].get('region')}"
        )

    async def get_droplet_create_info(self):
        mapping = {
            "nyc3": {"name": "rhino-proxy-snapshot-nyc3", "id": 57943519},
            "sfo2": {"name": "rhino-proxy-snapshot-sfo2", "id": 57943358},
            # "tor1": {"name": "rhino-proxy-snapshot-tor1", "id": 57943502},
        }

        return random.choice(list(mapping.items()))

    async def get_random_droplet(self):
        await self.ensure_min_droplets()

        droplet_id, droplet_blob = random.choice(list(self.droplets.items()))
        return droplet_id

    async def affix_ip(self, ytdl_blob, droplet_id):
        ytdl_blob["proxy"] = f"{self.droplets[droplet_id].get('ip_address')}:9049"
        return ytdl_blob

    async def get_droplet_info(self, droplet_id):
        return self.droplets[droplet_id]

    async def ensure_min_droplets(self):
        if len(self.droplets) < self.min_droplets:
            for _ in range(len(self.droplets), self.min_droplets):
                await self.create_droplet()