diff --git a/main.py b/main.py
index e59a039..37ec247 100644
--- a/main.py
+++ b/main.py
@@ -43,27 +43,11 @@ def now(receiver_email: str):
     generate_email(receiver_email)
     
 def generate_email(receiver_email: str):
-    #services.apt.require_root()
-    #if not services.apt.detect_apt():
-    #    print("Apt not found on this system.")
-    #    exit(1)
-    #updates = services.apt.check_updates()
-    
-    # TODO: Remove this after testing
-    updates = {
-        "openssl": {
-            "security": True,
-            "installed_version": "1.0.0",
-            "latest_version": "1.1.0",
-            "repo": "main",
-        },
-        "vim": {
-            "security": False,
-            "installed_version": "2.0.0",
-            "latest_version": "2.1.0",
-            "repo": "universe",
-        }
-    }
+    services.apt.require_root()
+    if not services.apt.detect_apt():
+        print("Apt not found on this system.")
+        exit(1)
+    updates = services.apt.check_updates()
     
     # Get how many security updates are available
     security_updates = 0
@@ -76,8 +60,9 @@ def generate_email(receiver_email: str):
     # Get how many general updates are available
     general_updates = total_updates - security_updates
     
-    if general_updates == 0:
-        print("No general updates available.")
+    # Check if there are any updates at all
+    if total_updates == 0:
+        print("No updates available.")
         return
     
     # Get system hostname
@@ -86,16 +71,12 @@ def generate_email(receiver_email: str):
     # Create email subject
     subject = f"{hostname}> {security_updates} security updates, {general_updates} general updates available"
     
-    chunk = f'' # define chunk
-    
-    # Iterate over each security package and create a list of updates in html
+    # Build security updates section
     security_chunks = ""
     
     for package in updates:
         if not updates[package]["security"]:
             continue
-        if security_updates == 0:
-            continue
         chunk = f'''
     <tr>
         <td style="border: 1px solid #ddd; padding: 8px;">{package}</td>
@@ -104,10 +85,11 @@ def generate_email(receiver_email: str):
         <td style="border: 1px solid #ddd; padding: 8px;">{updates[package]["repo"]}</td>
     </tr>
         '''
+        security_chunks += chunk
     
-    security_chunks += chunk
-    
-    security = f'''
+    security = ""
+    if security_updates > 0:
+        security = f'''
 <h1 style="font-family: Arial, sans-serif;">Security</h1>
 <table style="border-collapse: collapse; width: 100%; font-family: Arial, sans-serif;">
     <tr>
@@ -120,15 +102,12 @@ def generate_email(receiver_email: str):
 </table>
 '''
 
-    # Iterate over each general package and create a list of updates in html
-    
+    # Build general updates section
     general_chunks = ""
     
     for package in updates:
         if updates[package]["security"]:
             continue
-        if general_updates == 0:
-            continue
         chunk = f'''
     <tr>
         <td style="border: 1px solid #ddd; padding: 8px;">{package}</td>
@@ -137,10 +116,11 @@ def generate_email(receiver_email: str):
         <td style="border: 1px solid #ddd; padding: 8px;">{updates[package]["repo"]}</td>
     </tr>
         '''
-        
-    general_chunks += chunk
+        general_chunks += chunk
     
-    general = f'''
+    general = ""
+    if general_updates > 0:
+        general = f'''
 <h1 style="font-family: Arial, sans-serif;">General</h1>
 <table style="border-collapse: collapse; width: 100%; font-family: Arial, sans-serif;">
     <tr>
@@ -153,9 +133,7 @@ def generate_email(receiver_email: str):
 </table>
 '''
     
-    html = ""
-    html += security
-    html += general
+    html = security + general
     
     services.mail.send_email(receiver_email, subject, html)