forked from bellwether/minerva
91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
import boto3
|
|
import random
|
|
import os
|
|
import stat
|
|
import pprint
|
|
|
|
from minerva.machine import Machine
|
|
from minerva.cluster import Cluster
|
|
|
|
from fabric import Connection
|
|
|
|
pp = pprint.PrettyPrinter(indent=4)
|
|
|
|
# Used for interacting with AWS
|
|
class Pier:
|
|
def __init__(self,
|
|
handler = None,
|
|
subnet_id = None,
|
|
sg_groups = [],
|
|
key_pair = None, # (keypair name, keypair privkey pair)
|
|
iam = None,
|
|
):
|
|
self.handler = handler
|
|
self.session = handler.session
|
|
self.ec2 = self.session.client("ec2")
|
|
self.subnet_id = subnet_id
|
|
self.groups = sg_groups
|
|
self.iam = iam
|
|
self.account = self.session.client("sts").get_caller_identity()['Account']
|
|
self.ecr = f"{self.account}.dkr.ecr.{self.session.region_name}.amazonaws.com"
|
|
|
|
if key_pair:
|
|
print(f"Using keypair {key_pair}")
|
|
self.key_pair_name = key_pair[0]
|
|
self.key_path = os.path.expanduser(key_pair[1])
|
|
else:
|
|
self.make_key_pair()
|
|
|
|
|
|
def make_key_pair(self):
|
|
self.key_pair_name = f"Minerva-{random.random()}"
|
|
print(f"making keypair ({self.key_pair_name})")
|
|
self.key = self.ec2.create_key_pair(KeyName=self.key_pair_name)
|
|
|
|
# TODO add some randomization here
|
|
self.key_path = "/tmp/key"
|
|
with open(self.key_path, "w") as f:
|
|
f.write(self.key['KeyMaterial'])
|
|
os.chmod(self.key_path, stat.S_IRUSR | stat.S_IWUSR)
|
|
|
|
|
|
def machine(self, **kwargs):
|
|
return Machine(self, **kwargs)
|
|
|
|
|
|
# TODO make this fetch the instance id from the IP (filtered from the list
|
|
# of all instances)
|
|
def from_ip(self, ip, username, key_path=None, name=None):
|
|
mach = Machine(self, username=username)
|
|
|
|
mach.username = username
|
|
mach.public_ip = ip
|
|
mach.ip = ip
|
|
|
|
if key_path:
|
|
mach.key_pair = ("", key_path)
|
|
else:
|
|
mach.key_pair = (self.key_pair_name, self.key_path)
|
|
|
|
mach.ssh = Connection(ip,
|
|
username,
|
|
connect_kwargs = {
|
|
"key_filename": mach.key_pair[1]
|
|
}
|
|
)
|
|
mach.ssh.open()
|
|
|
|
return mach
|
|
|
|
|
|
def t3_med(self, num):
|
|
return self.machine(ami = "ami-0a538467cc9da9bb2",
|
|
instance_type = "t3.medium",
|
|
username = "ubuntu",
|
|
name = f"blah-{num}",
|
|
variables = {"num": num})
|
|
|
|
|
|
def cluster(self, *args, **kwargs):
|
|
return Cluster(self, *args, **kwargs)
|
|
|