Last active
June 7, 2020 09:45
-
-
Save srhinos/b2a831004294dbb778150680ca700422 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment