diff --git a/test.py b/test.py new file mode 100755 index 0000000..a2fd7f2 --- /dev/null +++ b/test.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 + +import argparse +import datetime +import tempfile +import subprocess +import re +import json +import time +import socket +import os.path +import urllib.request +import urllib.error + +scriptdir = os.path.dirname(os.path.realpath(__file__)) + +def log(*x): + print(datetime.datetime.now(), ':', *x) + +def doapi(token, path, args_or_method=None): + if args_or_method is None: + method = 'GET' + data = None + elif isinstance(args_or_method, dict): + method = 'POST' + data = json.dumps(args_or_method).encode() + else: + method = args_or_method + data = None + req = urllib.request.Request( + url='https://api.digitalocean.com' + path, + method=method, + data=data, + headers={ + 'Authorization': 'Bearer {}'.format(token), + 'Content-Type': 'application/json'}) + with urllib.request.urlopen(req) as hnd: + res = hnd.read().decode() + if res: + return json.loads(res) + else: + return None + +def waitssh(ipaddr): + log('Waiting for SSH ...') + sshup = False + while not sshup: + try: + sock = socket.socket() + sock.settimeout(5) + sock.connect((ipaddr, 22)) + sshup = True + except OSError: + pass + finally: + sock.close() + time.sleep(5) + +def main(args): + # parse args + parser = argparse.ArgumentParser(description='tests the install script') + parser.add_argument('access_token', help='DigitalOcean API access token') + args = parser.parse_args(args) + token = args.access_token + # check the token + log('Validating access token ...') + try: + doapi(token, '/v2/droplets') + except urllib.error.URLError: + log('Token is invalid.') + raise SystemExit + # create tempdir + tempdir = tempfile.TemporaryDirectory() + tempname = re.sub(r'[^A-Za-z0-9]+', '-', tempdir.name) + # generate a key + log('Generating SSH key ...') + subprocess.check_call(['ssh-keygen', '-f', tempdir.name + '/key', '-N', '', '-q']) + # add the key + log('Adding key ...') + keyid = doapi(token, '/v2/account/keys', { + 'name': tempname, + 'public_key': open(tempdir.name + '/key.pub').read() + })['ssh_key']['id'] + # create new droplet + log('Creating droplet ...') + dropid = doapi(token, '/v2/droplets', { + 'name': tempname, + 'region': 'sfo1', + 'size': '512mb', + 'image': 'debian-7-0-x64', + 'ssh_keys': [keyid] + })['droplet']['id'] + # wait for completion + completed = False + while not completed: + time.sleep(5) + result = doapi(token, '/v2/droplets/{}/actions'.format(dropid)) + completed = all(x['status'] == 'completed' for x in result['actions']) + # get ip address + v4nets = doapi(token, '/v2/droplets/{}'.format(dropid))['droplet']['networks']['v4'] + ipaddr = next(x['ip_address'] for x in v4nets if x['type'] == 'public') + # wait for SSH to start + waitssh(ipaddr) + # run the script + commonsshargs = [ + 'ssh', + '-i', tempdir.name + '/key', + '-o', 'StrictHostKeyChecking=no', + 'root@{}'.format(ipaddr)] + log('Running script ...') + with open(scriptdir + '/install.sh', 'rb') as scr: + subprocess.check_call( + commonsshargs + ['cat > install.sh && yes "" | bash install.sh'], + stdin=scr) + # wait for SSH to restart + time.sleep(5) + waitssh(ipaddr) + # check for /etc/arch-release + try: + subprocess.check_call(commonsshargs + ['cat /etc/arch-release']) + success = True + except subprocess.CalledProcessError: + success = False + # cleanup + if success: + log('>>> SUCCESS! Deleting droplet... <<<') + doapi(token, '/v2/droplets/{}'.format(dropid), 'DELETE') + else: + log('>>> FAILURE! Droplet saved for analysis. Remember to delete it to avoid charges. <<<') + doapi(token, '/v2/account/keys/{}'.format(keyid), 'DELETE') + +if __name__ == '__main__': + import sys + try: + main(sys.argv[1:]) + except SystemExit: + sys.exit(1)