From b0ebb5d503a21130c0d2c525c8c7c7218c32087d Mon Sep 17 00:00:00 2001
From: Vinicius Silva <vinicius@pop-os.localdomain>
Date: Tue, 16 May 2023 14:53:41 -0300
Subject: [PATCH] [Fix] Exceptions not captured

---
 danix/db/models.py | 32 ++++++++++++++++----------------
 danix/main.py      | 25 ++++++++++++++++---------
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/danix/db/models.py b/danix/db/models.py
index 686d103..e41eeb3 100644
--- a/danix/db/models.py
+++ b/danix/db/models.py
@@ -4,7 +4,6 @@ from danixfs import Danix
 from datetime import datetime
 from settings import SNAPSHOT_LIMIT
 from django.core.exceptions import ValidationError
-
 class Environment(models.Model):
     filesystem_name = models.UUIDField()
     status         = models.BooleanField(default=True)
@@ -39,19 +38,21 @@ class Environment(models.Model):
 
     @staticmethod
     def rm_environment(filesystem_name):
+        try:
+            environment = Environment.objects.filter(filesystem_name=filesystem_name)
 
-        environment = Environment.objects.filter(filesystem_name=filesystem_name)
+            if environment.count() == 0:
+                print("Environment does not exist!")
+                exit(1)
+            else:
 
-        if environment.count() == 0:
+                Danix.rm(filesystem_name)
+                env = environment.first()
+                env.delete()
+                print(f"Environment removed successfully! - {filesystem_name}")
+        except Exception:
             print("Environment does not exist!")
             exit(1)
-        else:
-
-            Danix.rm(filesystem_name)
-            env = environment.first()
-            env.delete()
-            print("Environment removed successfully!")
-            exit(0)
 
 
     @staticmethod
@@ -138,8 +139,7 @@ class Snapshot(models.Model):
                 if resp == 0:
                     Snapshot.objects.get(snapshot_name=snapshot_name).delete()
 
-                    print("Snapshot removed successfully")
-                    exit(0)
+                    print(f"Snapshot removed successfully - {snapshot_name}")
                 else:
                     print("Error: Snapshot can not remove")
                     exit(1)
@@ -179,10 +179,11 @@ class Snapshot(models.Model):
     def create(subsystem_name):
         try:
             environment = Environment.objects.get(filesystem_name=subsystem_name)
+           
             environment_id = environment.id
             snapshots      = Snapshot.objects.filter(environment_id=environment_id)
 
-            if snapshots.count() > int(SNAPSHOT_LIMIT):
+            if snapshots.count() >= int(SNAPSHOT_LIMIT):
                 print("Snapshot limit exceeded! Please remove 1 snapshot to continue")
                 exit(1)
             else:
@@ -207,8 +208,7 @@ class Snapshot(models.Model):
                     print("Snapshot create error!")
                     exit(1)
 
-        except ValidationError:
-
+        except Exception:
             print("Snapshot create error: Environment does not exist!")
             exit(1)
 
@@ -225,7 +225,7 @@ class Snapshot(models.Model):
             for snapshot in snapshots:
 
                 name = str(snapshot.snapshot_name)
-                lastsnapshot_icon = "🟢 Yes" if snapshot.last else "🟠 No"
+                lastsnapshot_icon = "🟢 Yes" if snapshot.last else "🟠 No "
 
                 if snapshot.environment_id:
                     environment_name = Environment.objects.filter(id=snapshot.environment_id.id).first().filesystem_name
diff --git a/danix/main.py b/danix/main.py
index 45ee278..2677638 100644
--- a/danix/main.py
+++ b/danix/main.py
@@ -27,13 +27,6 @@ usages.add_argument("-sb", "--snapshotback",        help="Back snapshot",   requ
 
 usages.add_argument("-o", "--option",   choices=["clike", "java", "python", "ruby"], required=False)
 
-#fullstack_package = parser.add_argument_group("WebStack", "")
-
-#fullstack_package.add_argument("-d","--db", type=str, required=False)
-#fullstack_package.add_argument("-w", "--webserver", type=str, required=False)
-#fullstack_package.add_argument("-f", "--framework", type=str, required=False)
-#fullstack_package.add_argument("-e", "--env_name", type=str, required=False)
-
 args = parser.parse_args()
 
 languanges_and_softwares = {
@@ -55,7 +48,15 @@ if args.start:
     Environment.start_environment(args.start)
 
 if args.rm:
-    Environment.rm_environment(args.rm)
+     user_confirm = input("Type 'y' to continue: ")
+     
+     if user_confirm == 'y':
+        environments = str(args.rm).split(" ")
+    
+        for environment in environments:
+            Environment.rm_environment(environment)
+            exit(0)
+     print("[Danix]: System abort!")
 
 if args.stop:
     Environment.stop_environment(args.stop)
@@ -84,5 +85,11 @@ if args.snapshotremove:
     user_confirm = input("Type y to continue: ")
     
     if user_confirm == 'y':
-        Snapshot.rm_snapshot(args.snapshotremove)
+
+        snapshots = str(args.snapshotremove).split(" ")
+    
+        for snapshot in snapshots:
+            Snapshot.rm_snapshot(snapshot)
+            exit(0)
+
     print("[Danix]: System abort!")
\ No newline at end of file