implement forbidden commands; prune vagrant boxes; improve status report; ref...

implement forbidden commands; prune vagrant boxes; improve status report; ref liberty/ldh_developer#8

box/command

@@ -27,6 +27,35 @@ XDG_PB_CONFIG_HOME = Path(os.environ['HOME'],'.config','ldh_developer','playbook
 VG_REQ_PLUGINS = {'libvirt': 'vagrant-libvirt', 'digital_ocean': 'vagrant-digitalocean'}
 HERE = Path(__file__).parent
 DO_SSH_KEY_PATH = Path('~/.ssh/shipwright/').expanduser()
+VAGRANT_CONFIG_HOME = Path(os.environ.get('VAGRANT_HOME', str(Path.home().joinpath('.vagrant.d'))))
+VAGRANT_DB_PATH = Path(VAGRANT_CONFIG_HOME, 'data', 'machine-index', 'index')
+
+
+class VagrantDB(object):
+    @property
+    def boxes(self):
+        data = json.loads(VAGRANT_DB_PATH.open().read())
+        vgt_boxes = data['machines']
+        return {
+            vgt_boxes[box_id]['name']:{
+                'uid': box_id,
+                'provider': vgt_boxes[box_id]['provider'],
+                'state': vgt_boxes[box_id]['state'],
+                'vagrantfile_path': vgt_boxes[box_id]['vagrantfile_path'],
+            } for box_id in vgt_boxes
+        }
+
+    def provider_for(self, hostname):
+        try:
+            return self.boxes[hostname]['provider']
+        except KeyError:
+            return None
+
+    @classmethod
+    def prune(cls):
+        print("Removing vagrant's invalid entries ...")
+        subprocess.run(['vagrant', 'global-status', '--prune'], capture_output=True)
+        print("Done")
 
 
 def ensure_plugins():
@@ -36,7 +65,7 @@ def ensure_plugins():
         universal_newlines=True
     )
     if cp.returncode != 0:
-        print("Ooops, there's be something wrong with vagrant!!?")
+        print("Ooops, there's something wrong with vagrant!!?")
         print(cp.stdout)
         print(cp.stderr)
         sys.exit(1)
@@ -48,6 +77,18 @@ def ensure_plugins():
         sys.exit(1)
 
 
+def forbdden_for_digital_ocean(func):
+    def wrapper(*args, **kwargs):
+        vg_db = VagrantDB()
+        hostname = args[0]
+        if vg_db.provider_for(hostname) == 'digital_ocean':
+            print("Error: This command is forbidden for the digital_ocean provider.")
+            exit(1)
+        else:
+            func(*args, **kwargs)
+    return wrapper
+
+
 def only_if_box_exist(func):
     def wrapper(*args, **kwargs):
         hostname = args[0]
@@ -60,6 +101,13 @@ def only_if_box_exist(func):
     return wrapper
 
 
+def prune_after_call(func):
+    def wrapper(*args, **kwargs):
+        func(*args, **kwargs)
+        VagrantDB.prune()
+    return wrapper
+
+
 def is_ipaddress(ip_addr):
     try:
         ipaddress.ip_address(ip_addr)
@@ -73,15 +121,13 @@ def print_status():
     Prints the status of all the boxes managed by Shipwright
     """
     box_paths = [str(d) for d in XDG_BOX_CONFIG_HOME.iterdir() if d.is_dir()]
-    vgt_db = Path(
-        os.environ.get('VAGRANT_HOME', str(Path.home().joinpath('.vagrant.d'))),
-        'data', 'machine-index', 'index'
-    )
     count = 0
-    if vgt_db.exists():
-        print('hostname \t State\t IP Address \t Provider')
-        db = json.loads(vgt_db.open().read())
-        for k, v in db['machines'].items():
+    if VAGRANT_CONFIG_HOME.exists():
+        print('-----------------------------------------------------------------')
+        print('| hostname       | State     | Provider       | IP Address      |')
+        print('-----------------------------------------------------------------')
+        db = VagrantDB()
+        for hostname, v in db.boxes.items():
             vagrant_wd = v['vagrantfile_path']
             if vagrant_wd in box_paths:
                 # get the ip address
@@ -97,11 +143,13 @@ def print_status():
                         if is_ipaddress(str(fld)):
                             ip_addr = str(fld)
                             break
-                print(f"{v['name']} \t {v['state']}\t {ip_addr}")
+                print(f"| {hostname:15}| {v['state']:10}| {v['provider']:15}| {ip_addr:15} |")
                 count+=1
-    print('Total {} box(es)'.format(count))
+    print('-----------------------------------------------------------------')
+    print(f'Total {count} box(es)')
 
 
+@prune_after_call
 def box_create(hostname, ram, cpus, provider, token):
     """
     This creates a box using the template
@@ -142,16 +190,21 @@ def box_create(hostname, ram, cpus, provider, token):
     subprocess.run(['vagrant','up'], cwd=vagrant_wd)
 
 
+@prune_after_call
+@forbdden_for_digital_ocean
 @only_if_box_exist
 def box_destroy(hostname):
     """
     Destroy a box, or all
     """
     vagrant_wd = XDG_BOX_CONFIG_HOME.joinpath(hostname)
-    subprocess.run(['vagrant','destroy'], cwd=vagrant_wd)
-    shutil.rmtree(vagrant_wd)
+    r = subprocess.run(['vagrant','destroy'], cwd=vagrant_wd)
+    if r.returncode == 0:
+        shutil.rmtree(vagrant_wd)
 
 
+@prune_after_call
+@forbdden_for_digital_ocean
 @only_if_box_exist
 def run_playbook(hostname, playbook, retry_file=None):
     """
@@ -171,6 +224,7 @@ def run_playbook(hostname, playbook, retry_file=None):
     subprocess.run(_args, cwd=vagrant_wd)
 
 
+@forbdden_for_digital_ocean
 @only_if_box_exist
 def box_start(hostname):
     """
@@ -180,6 +234,8 @@ def box_start(hostname):
     subprocess.run(['vagrant','up'], cwd=vagrant_wd)
 
 
+@prune_after_call
+@forbdden_for_digital_ocean
 @only_if_box_exist
 def box_halt(hostname):
     """
@@ -190,7 +246,7 @@ def box_halt(hostname):
 
 
 @only_if_box_exist
-def invoke_shell(hostname):
+def box_ssh_connect(hostname):
     """
     Open a SSH session to the hostname
     """
@@ -199,7 +255,7 @@ def invoke_shell(hostname):
 
 
 @only_if_box_exist
-def which(hostname):
+def box_get_working_directory(hostname):
     """
     Print the directory of a box by it's hostname
     """
@@ -281,9 +337,9 @@ if __name__ == '__main__':
     elif cmd_args.command == 'destroy':
         box_destroy(sub_args.hostname)
     elif cmd_args.command == 'ssh':
-        invoke_shell(sub_args.hostname)
+        box_ssh_connect(sub_args.hostname)
     elif cmd_args.command == 'which':
-        which(sub_args.hostname)
+        box_get_working_directory(sub_args.hostname)
     elif cmd_args.command == 'open':
         box_gio_open(sub_args.hostname)
     else: