fixed user data, added pub ips, colors

Author: chadpierce

dew.py

@@ -11,13 +11,16 @@
 #   list available images: python3 dew.py -i
 #   show help: python3 dew.py -h
 #
-# TODO user data is iffy and needs testing
 ##################################################
 import sys
 import time
 import requests
 import json
-
+cRed = '\033[31m'
+cOrange = '\033[33m'
+cPurple = '\033[35m'
+cCyan = '\033[36m'
+cReset = '\033[0m'
 # api_token - accounts > api > generate new token
 # ssh_keys - settings > security > SSH Keys
 # tags - use to automatically assign firewalls, etc
@@ -29,13 +32,12 @@
 droplet_name = 'pwnbox'  # default if you leave input blank
 droplet_size = 's-1vcpu-1gb'  #cheapest $5 droplet
 droplet_image = 'debian-10-x64'  # 'ubuntu-20-04-x64'
-tag = 'pwn'
+tag = 'tagname'  # for more than one tag use the format: tag = 'tag1","tag2' maybe i'll fix this
 region = 'tor1'
 user_data = ''  # this is populated in create_droplet()
 auto_script = '''#!/bin/bash
-
 apt-get -y update
-apt-get -y install nmap wget curl tmux
+apt-get -y install nmap wget curl tmux htop
 touch /root/donezo
 '''
 # ^^^ change that stuff ^^^
@@ -45,77 +47,96 @@
     'Authorization': 'Bearer ' + api_token,
 }
 
-data = '{ \
-    "name":"' + droplet_name + '", \
-    "region":"' + region  + '", \
-    "size":"' + droplet_size  + '", \
-    "image":"' + droplet_image + '", \
-    "ssh_keys":["' + ssh_key_fingerprint + '"], \
-    "backups":false, \
-    "ipv6":false, \
-    "user_data":"' + user_data + '", \
-    "private_networking":null, \
-    "volumes": null, \
-    "tags":["' + tag + '"] \
-}'
-
 
 def create_droplet():
+    print()
     global droplet_name
     name_drop = input('enter droplet name (default is ' + droplet_name + ') => ')
     if name_drop != '':
         droplet_name = name_drop
     do_auto = input('include automated script (may not work)? (y or n) => ')
-    if do_auto == 'n':
-        user_data = ''
-    elif do_auto == 'y':
+    if do_auto == 'y':
         user_data = auto_script
+    elif do_auto == 'n':
+        user_data = ''
     else:
         print('that was not y or n')
         sys.exit()
+    data = '{ \
+        "name":"' + droplet_name + '", \
+        "region":"' + region  + '", \
+        "size":"' + droplet_size  + '", \
+        "image":"' + droplet_image + '", \
+        "ssh_keys":["' + ssh_key_fingerprint + '"], \
+        "backups":false, \
+        "ipv6":false, \
+        "user_data":"' + user_data + '", \
+        "private_networking":null, \
+        "volumes": null, \
+        "tags":["' + tag + '"] \
+    }'
+
     response = requests.post('https://api.digitalocean.com/v2/droplets', headers=headers, data=data)
     jresp = response.json()
     droplet_id = jresp['droplet']['id']
     print('will generate droplet then sleep before getting IP...')
     print()
-    print('droplet id: ' + str(droplet_id))
+    print('droplet id: ' + cCyan + str(droplet_id) + cReset)
     time.sleep(10)
     response = requests.get('https://api.digitalocean.com/v2/droplets/' + str(droplet_id) + '/', headers=headers)
     jresp = response.json()
-    ip1 = jresp['droplet']['networks']['v4'][0]['ip_address']  # TODO make this better
-    iptype1 = jresp['droplet']['networks']['v4'][0]['type']
-    ip2 = jresp['droplet']['networks']['v4'][1]['ip_address']
-    iptype2 = jresp['droplet']['networks']['v4'][1]['type']
-    if iptype1 == 'public':
-        print('public ip: ' + str(ip1))
-    elif iptype2 == 'public':
-        print('public ip: ' + str(ip2))
-    else:
-        print('ERROR: no public ip? printing all ip4 info')
+    pub_ip = 'none' 
+    priv_ip = 'none' 
+    for i in range(len(jresp['droplet']['networks']['v4'])):
+        if jresp['droplet']['networks']['v4'][i]['type'] == 'public':
+            pub_ip = jresp['droplet']['networks']['v4'][i]['ip_address']
+        if jresp['droplet']['networks']['v4'][i]['type'] == 'private':
+            priv_ip = jresp['droplet']['networks']['v4'][i]['ip_address']    
+    if pub_ip == '' or priv_ip == '':
+        print(cRed + 'ERROR: ' + cReset + 'ip error, printing all ip info:')
         print(jresp['droplet']['networks']['v4'])
+    else:
+        print('public ip: ' + cOrange + pub_ip + cReset)
+        print('private ip: ' + cPurple + priv_ip + cReset)
+        print('tags: ' + str(jresp['droplet']['tags']))
+    print()
 
 
 def del_droplet(drop_id):
+    print()
     if int(drop_id) in safe_droplets:
         print('you dont want to delete that one!')
     else:
         response = requests.delete('https://api.digitalocean.com/v2/droplets/' + drop_id, headers=headers)
         if response.status_code == 204:
-            print('droplet deleted')
+            print(cCyan + 'droplet deleted' + cReset)
         else:
-            print('ERROR: something went wrong!')
+            print(cRed + 'ERROR: ' + cReset + 'something went wrong!')
+    print()
 
 
 def list_droplets():
+    print()
     response = requests.get('https://api.digitalocean.com/v2/droplets', headers=headers)
     jresp = response.json()
-    ipaddr = 'error: public ip not found' 
-    print('name\t\tid\t\tip\n----------\t----------\t--------------')
+    print('name\t\tid\t\t\tpublic ip\t\tprivate ip\n----------\t-----------------\t--------------\t\t--------------')
     for d in jresp['droplets']:
+        if d['id'] in safe_droplets:
+            safe_color = cRed
+            safe_string = ' (safe)'
+        else:
+            safe_color = cCyan
+            safe_string = '\t'
+        pub_ip = 'none' 
+        priv_ip = 'none' 
         for i in range(len(d['networks']['v4'])):
             if d['networks']['v4'][i]['type'] == 'public':
-                ipaddr = d['networks']['v4'][i]['ip_address']
-        print(d['name'] + '\t\t' + str(d['id']) + '\t' + str(ipaddr))
+                pub_ip = d['networks']['v4'][i]['ip_address']
+            if d['networks']['v4'][i]['type'] == 'private':
+                priv_ip = d['networks']['v4'][i]['ip_address']
+        print(d['name'] + '\t\t' + safe_color + str(d['id']) + safe_string + cReset \
+            + '\t' + str(pub_ip) + '\t\t' + str(priv_ip))
+    print()
 
 
 def list_images():
@@ -128,6 +149,7 @@ def list_images():
     res = []  # remove duplicates 
     [res.append(x) for x in image_list if x not in res]
     print(*res, sep='\n')
+    print()
 
 
 def print_help():
@@ -141,6 +163,7 @@ def print_help():
       '''
     print(help_str)
 
+
 # start script        
 if len(sys.argv) > 1:
     if sys.argv[1] == '-d':
@@ -164,4 +187,5 @@ def print_help():
 elif len(sys.argv) == 1:
     print_help()
 else:
-    add_droplet()
+    print('something went wrong')
+    sys.exit()