From dd876a8b2cd2178b5d2f8e840b2fad17d28662ab Mon Sep 17 00:00:00 2001
From: Jordan Rodgers <com6056@gmail.com>
Date: Thu, 26 Apr 2018 02:59:09 -0400
Subject: [PATCH] tear down SSH tunnel every time console button is clicked and
 on suspend, yapf a few things

---
 proxstar/__init__.py | 28 ++++++----------------------
 proxstar/mail.py     | 15 ++++++++++-----
 proxstar/tasks.py    |  4 +++-
 3 files changed, 19 insertions(+), 28 deletions(-)

diff --git a/proxstar/__init__.py b/proxstar/__init__.py
index e15d043..b9fdbfa 100644
--- a/proxstar/__init__.py
+++ b/proxstar/__init__.py
@@ -215,6 +215,7 @@ def vm_power(vmid, action):
             vm.reset()
         elif action == 'suspend':
             vm.suspend()
+            send_stop_ssh_tunnel(vmid)
         elif action == 'resume':
             vm.resume()
         return '', 200
@@ -238,31 +239,14 @@ def vm_console(vmid):
     proxmox = connect_proxmox()
     if user.rtp or int(vmid) in user.allowed_vms:
         vm = VM(vmid)
+        stop_ssh_tunnel(vm.id, ssh_tunnels)
         port = str(5900 + int(vmid))
         token = add_vnc_target(port)
         node = "{}.csh.rit.edu".format(vm.node)
-        tunnel = next((tunnel for tunnel in ssh_tunnels
-                       if tunnel.local_bind_port == int(port)), None)
-        if tunnel:
-            if tunnel.ssh_host != node:
-                print(
-                    "Tunnel already exists for VM {} to the wrong Proxmox node.".
-                    format(vmid))
-                tunnel.stop()
-                ssh_tunnels.remove(tunnel)
-                print("Creating SSH tunnel to {} for VM {}.".format(
-                    node, vmid))
-                tunnel = start_ssh_tunnel(node, port)
-                ssh_tunnels.append(tunnel)
-                vm.start_vnc(port)
-            else:
-                print("Tunnel already exists to {} for VM {}.".format(
-                    node, vmid))
-        else:
-            print("Creating SSH tunnel to {} for VM {}.".format(node, vmid))
-            tunnel = start_ssh_tunnel(node, port)
-            ssh_tunnels.append(tunnel)
-            vm.start_vnc(port)
+        print("Creating SSH tunnel to {} for VM {}.".format(node, vm.id))
+        tunnel = start_ssh_tunnel(node, port)
+        ssh_tunnels.append(tunnel)
+        vm.start_vnc(port)
         return token, 200
     else:
         return '', 403
diff --git a/proxstar/mail.py b/proxstar/mail.py
index 049047d..5df597a 100644
--- a/proxstar/mail.py
+++ b/proxstar/mail.py
@@ -24,16 +24,19 @@ def send_vm_expire_email(user, vms):
     body = "The following VMs in Proxstar are expiring soon or have already expired:\n\n"
     for vm in vms:
         if vm[2] == -6:
-            body += "    - {} ({}) has expired (VM has been stopped and will be deleted in 1 day)\n".format(vm[1], vm[0])
+            body += "    - {} ({}) has expired (VM has been stopped and will be deleted in 1 day)\n".format(
+                vm[1], vm[0])
         elif vm[2] < 0:
             body += "    - {} ({}) has expired (VM has been stopped and will be deleted in {} days)\n".format(
                 vm[1], vm[0], (7 + int(vm[2])))
         elif vm[2] == 0:
-            body += "    - {} ({}) expires today (VM has been stopped and will be deleted in 7 days)\n".format(vm[1], vm[0])
+            body += "    - {} ({}) expires today (VM has been stopped and will be deleted in 7 days)\n".format(
+                vm[1], vm[0])
         elif vm[2] == 1:
             body += "    - {} ({}) expires in 1 day\n".format(vm[1], vm[0])
         else:
-            body += "    - {} ({}) expires in {} days\n".format(vm[1], vm[0], vm[2])
+            body += "    - {} ({}) expires in {} days\n".format(
+                vm[1], vm[0], vm[2])
     body += "\nPlease login to Proxstar (https://proxstar.csh.rit.edu/) and renew any VMs you would like to keep."
     send_email(toaddr, subject, body)
 
@@ -44,8 +47,10 @@ def send_rtp_vm_delete_email(vms):
     body = "The following VMs in Proxstar have expired and will be deleted soon:\n\n"
     for vm in vms:
         if vm[2] == -6:
-            body += "    - {} ({}) will be deleted in 1 day\n".format(vm[1], vm[0])
+            body += "    - {} ({}) will be deleted in 1 day\n".format(
+                vm[1], vm[0])
         else:
-            body += "    - {} ({}) will be deleted in {} days\n".format(vm[1], vm[0], (7 + int(vm[2])))
+            body += "    - {} ({}) will be deleted in {} days\n".format(
+                vm[1], vm[0], (7 + int(vm[2])))
     body += "\nPlease verify this list to ensure there aren't any pools included in Proxstar that shouldn't be."
     send_email(toaddr, subject, body)
diff --git a/proxstar/tasks.py b/proxstar/tasks.py
index 9320b75..903c438 100644
--- a/proxstar/tasks.py
+++ b/proxstar/tasks.py
@@ -95,7 +95,9 @@ def process_expiring_vms_task():
                     if days <= 0:
                         vm.stop()
                 elif days == -7:
-                    print("Deleting {} ({}) as it has been a week since expiration.".format(vm.name, vm.id))
+                    print(
+                        "Deleting {} ({}) as it has been a week since expiration.".
+                        format(vm.name, vm.id))
                     send_stop_ssh_tunnel(vm.id)
                     delete_vm_task(vm.id)
             if expiring_vms: