404 fix

app.py

@@ -137,7 +137,30 @@ def inject_subdomain():
 
 @app.errorhandler(404)
 def page_not_found(_):
-    return render_template("404.html", domain=request.host), 404
+    host = request.host
+    server_name = app.config["SERVER_NAME"]
+    server_parts = server_name.split('.')
+    host_parts = host.split('.')
+    show_domain = False
+
+    # Case 1: Direct match of main domain
+    if host == server_name:
+        show_domain = True
+    else:
+        # Extract potential subdomain
+        if host_parts[-len(server_parts):] == server_parts:
+            subdomain = '.'.join(host_parts[:-len(server_parts)])
+            
+            # Case 2: Subdomain doesn't exist and isn't reserved
+            if subdomain and subdomain not in RESERVED_SUBDOMAINS:
+                if not Site.query.filter_by(subdomain=subdomain).first():
+                    show_domain = True
+
+    return render_template(
+        "404.html", 
+        domain=host if show_domain else None,
+        is_main_domain=(host == server_name)
+    ), 404
 
 
 # Auth setup

templates/404.html

@@ -65,14 +65,25 @@
 	</div>
 
 	{% if current_user.is_authenticated %}
-	<a href="{{ url_for('dashboard') }}" class="btn cta-button">
+	{% if domain %}
-		🚀 Create {{ domain }}
+	<a href="{{ url_for('dashboard', _external=True) }}" class="btn cta-button">
+		🚀 Create {{ domain.split('.')[0]|default("Site", true) }}
 	</a>
 	{% else %}
 	<div class="auth-buttons">
+		{% if is_main_domain %}
 		<p>Start your web hosting journey today!</p>
-		<a href="{{ url_for('register') }}" class="btn">Sign Up</a>
+		<a href="{{ url_for('register', _external=True) }}">Sign Up</a>
-		<a href="{{ url_for('login') }}" class="btn">Login</a>
+		{% else %}
+		<p>The page you requested doesn't exist on this site.</p>
+		{% endif %}
+	</div>
+	{% endif %}
+	{% else %}
+	<div class="auth-buttons">
+		<p>Start your web hosting journey today!</p>
+		<a href="{{ url_for('login', _external=True) }}">Login</a>
+		<a href="{{ url_for('register', _external=True) }}">Sign Up</a>
 	</div>
 	{% endif %}