wip

app/routes/api.py

@@ -8,8 +8,8 @@ from flask import (
     redirect,
     url_for,
 )
-from flask_login import login_required
-from app.core.models import Subnet, Server, App, Port
+from flask_login import login_required, current_user
+from app.core.models import Subnet, Server, App, Port, Location
 from app.core.extensions import db
 from app.scripts.ip_scanner import scan
 import random
@@ -20,33 +20,51 @@ from datetime import datetime
 from flask import flash
 from app.utils.app_utils import is_port_in_use, validate_port_data
 from difflib import SequenceMatcher
+import json
 
 bp = Blueprint("api", __name__, url_prefix="/api")
 csrf = CSRFProtect()
 
 
 @bp.route("/subnets", methods=["GET"])
+@csrf.exempt
+@login_required
 def get_subnets():
     """Get all subnets grouped by site"""
-    subnets = Subnet.query.all()
-
-    # Group subnets by location (site)
-    sites = {}
-    for subnet in subnets:
-        location = subnet.location
-        if location not in sites:
-            sites[location] = []
-
-        sites[location].append(
-            {"id": subnet.id, "cidr": subnet.cidr, "location": location}
-        )
-
-    # Convert to list of site objects
-    result = [
-        {"name": site_name, "subnets": subnets} for site_name, subnets in sites.items()
-    ]
-
-    return jsonify(result)
+    try:
+        subnets = Subnet.query.filter_by(user_id=current_user.id).all()
+        
+        # Group subnets by location (site)
+        sites = {}
+        for subnet in subnets:
+            location = subnet.location_ref  # Make sure this attribute matches your model relationship
+            if not location:
+                location_name = "Unassigned"
+                location_id = None
+            else:
+                location_name = location.name
+                location_id = location.id
+            
+            if location_id not in sites:
+                sites[location_id] = {
+                    "name": location_name,
+                    "id": location_id,
+                    "subnets": []
+                }
+            
+            sites[location_id]["subnets"].append({
+                "id": subnet.id, 
+                "cidr": subnet.cidr, 
+                "location_id": location_id
+            })
+        
+        # Convert to list of site objects
+        result = list(sites.values())
+        
+        return jsonify(result)
+    except Exception as e:
+        print(f"Error loading subnets: {e}")  # Add debugging
+        return jsonify({"error": str(e)}), 500
 
 
 @bp.route("/subnets/<int:subnet_id>", methods=["GET"])
@@ -288,40 +306,60 @@ def add_app_port(app_id):
     """Add a port to an application"""
     app = App.query.get_or_404(app_id)
     
+    # Check if the user owns this app
+    if app.user_id != current_user.id:
+        flash("Unauthorized access", "danger")
+        return redirect(url_for("dashboard.app_view", app_id=app_id))
+    
     # Get port details from the form
     port_number = request.form.get("port_number")
     protocol = request.form.get("protocol", "TCP")
     description = request.form.get("description", "")
     
     # Validate the port
-    valid, clean_port, error = validate_port_data ( # validate_port_data(
+    valid, clean_port, error = validate_port_data(
         port_number,
-        protocol,
-        description,
-        app.server_id,
-        app_id
+        descriptions=description,
+        server_id=app.server_id,
+        exclude_app_id=app_id,
+        protocol=protocol
     )
     
     if not valid:
         flash(error, "danger")
         return redirect(url_for("dashboard.app_view", app_id=app_id))
     
-    # Create the new port
     try:
-        new_port = Port(
+        # Check if the port already exists for this app
+        existing_port = Port.query.filter_by(
             app_id=app_id,
             port_number=clean_port,
-            protocol=protocol,
-            description=description
-        )
-        db.session.add(new_port)
+            protocol=protocol
+        ).first()
+        
+        if existing_port:
+            # Update the existing port description
+            existing_port.description = description
+            flash(f"Port {clean_port}/{protocol} updated", "success")
+        else:
+            # Create a new port
+            port = Port(
+                port_number=clean_port,
+                protocol=protocol,
+                description=description,
+                app_id=app_id
+            )
+            db.session.add(port)
+            
+            flash(f"Port {clean_port}/{protocol} added successfully", "success")
+        
         db.session.commit()
-        flash(f"Port {clean_port}/{protocol} added successfully", "success")
+        return redirect(url_for("dashboard.app_view", app_id=app_id))
+    
     except Exception as e:
         db.session.rollback()
         flash(f"Error adding port: {str(e)}", "danger")
-    
-    return redirect(url_for("dashboard.app_view", app_id=app_id))
+        return redirect(url_for("dashboard.app_view", app_id=app_id))
 
 
 @bp.route("/app/<int:app_id>/ports", methods=["GET"])
@@ -337,7 +375,7 @@ def get_app_ports(app_id):
         "ports": [
             {
                 "id": port.id,
-                "number": port.number,
+                "port_number": port.port_number,
                 "protocol": port.protocol,
                 "description": port.description,
             }
@@ -396,7 +434,7 @@ def get_server_ports(server_id):
 
     # Get all ports associated with this server
     ports = Port.query.filter_by(server_id=server_id).all()
-    used_ports = [port.number for port in ports]
+    used_ports = [port.port_number for port in ports]
 
     return jsonify({"server_id": server_id, "used_ports": used_ports})
 
@@ -409,7 +447,7 @@ def get_free_port(server_id):
 
     # Get all ports associated with this server
     used_ports = [
-        port.number for port in Port.query.filter_by(server_id=server_id).all()
+        port.port_number for port in Port.query.filter_by(server_id=server_id).all()
     ]
 
     # Find the first free port (starting from 8000)
@@ -524,3 +562,65 @@ def validate_app_port():
         
     except ValueError:
         return jsonify({"valid": False, "message": "Invalid port number"})
+
+
+@bp.route("/locations", methods=["POST"])
+@login_required
+def create_location():
+    """API endpoint to create a new location"""
+    data = request.json
+    
+    if not data or not data.get('name'):
+        return jsonify({'error': 'Location name is required'}), 400
+        
+    try:
+        location = Location(
+            name=data.get('name'),
+            description=data.get('description', ''),
+            user_id=current_user.id
+        )
+        
+        db.session.add(location)
+        db.session.commit()
+        
+        return jsonify({
+            'id': location.id,
+            'name': location.name,
+            'description': location.description
+        }), 201
+        
+    except Exception as e:
+        db.session.rollback()
+        return jsonify({'error': str(e)}), 500
+
+
+@bp.route("/subnets", methods=["POST"])
+@login_required
+def create_subnet():
+    """API endpoint to create a new subnet"""
+    data = request.json
+    
+    if not data or not data.get('cidr') or not data.get('location_id'):
+        return jsonify({'error': 'CIDR and location_id are required'}), 400
+        
+    try:
+        subnet = Subnet(
+            cidr=data.get('cidr'),
+            location_id=data.get('location_id'),
+            user_id=current_user.id,
+            auto_scan=data.get('auto_scan', False),
+            active_hosts=json.dumps([])
+        )
+        
+        db.session.add(subnet)
+        db.session.commit()
+        
+        return jsonify({
+            'id': subnet.id,
+            'cidr': subnet.cidr,
+            'location_id': subnet.location_id
+        }), 201
+        
+    except Exception as e:
+        db.session.rollback()
+        return jsonify({'error': str(e)}), 500