VPN/scripts/seed_data.py

#!/usr/bin/env python3
"""
Seed data script for VPN Access Server.

Creates sample users and MAC addresses for testing purposes.
"""

import sys
import os

# Add the access module to the Python path
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'access'))

from db import db
from utils import setup_logging


def create_sample_users():
    """Create sample users for testing."""

    sample_users = [
        {
            'username': 'john.doe',
            'employee_name': 'John Doe',
            'employee_email': 'john.doe@company.com',
            'session_limit': 28800  # 8 hours
        },
        {
            'username': 'jane.smith',
            'employee_name': 'Jane Smith',
            'employee_email': 'jane.smith@company.com',
            'session_limit': 14400  # 4 hours
        },
        {
            'username': 'admin',
            'employee_name': 'Administrator',
            'employee_email': 'admin@company.com',
            'session_limit': 86400  # 24 hours
        }
    ]

    sample_mac_addresses = {
        'john.doe': ['00:11:22:33:44:55', '00:11:22:33:44:56'],
        'jane.smith': ['aa:bb:cc:dd:ee:ff'],
        'admin': ['12:34:56:78:9a:bc', '12:34:56:78:9a:bd', '12:34:56:78:9a:be']
    }

    try:
        # Insert users
        for user_data in sample_users:
            # Check if user already exists
            existing_user = db.get_user_by_username(user_data['username'])
            if existing_user:
                logger.info(f"User '{user_data['username']}' already exists, skipping...")
                continue

            # Insert user
            insert_user_query = """
                INSERT INTO employees (username, employee_name, employee_email, session_limit, is_active)
                VALUES (%s, %s, %s, %s, TRUE)
            """

            db.execute_query(insert_user_query, (
                user_data['username'],
                user_data['employee_name'],
                user_data['employee_email'],
                user_data['session_limit']
            ))

            logger.info(f"Created user: {user_data['username']}")

            # Get the created user to get the ID
            created_user = db.get_user_by_username(user_data['username'])
            if not created_user:
                logger.error(f"Failed to retrieve created user: {user_data['username']}")
                continue

            user_id = created_user['id']

            # Add MAC addresses for this user
            mac_list = sample_mac_addresses.get(user_data['username'], [])
            for mac_address in mac_list:
                if db.add_mac_address(user_id, mac_address):
                    logger.info(f"Added MAC address {mac_address} for user {user_data['username']}")
                else:
                    logger.warning(f"Failed to add MAC address {mac_address} for user {user_data['username']}")

        return True

    except Exception as e:
        logger.error(f"Error creating sample users: {e}")
        return False


def display_created_data():
    """Display the created sample data."""
    try:
        # Get all users
        users_query = """
            SELECT u.id, u.username, u.employee_name, u.employee_email,
                   u.session_limit, u.is_active,
                   GROUP_CONCAT(m.mac ORDER BY m.created_at) as mac_addresses
            FROM employees u
            LEFT JOIN mac_addresses m ON u.id = m.employee_id
            WHERE u.is_active = TRUE
            GROUP BY u.id
            ORDER BY u.username
        """

        users = db.execute_query(users_query, fetch='all')

        if users:
            logger.info("\n" + "="*80)
            logger.info("CREATED SAMPLE DATA:")
            logger.info("="*80)

            for user in users:
                logger.info(f"""
User: {user['username']} (ID: {user['id']})
  Name: {user['employee_name']}
  Email: {user['employee_email']}
  Session Limit: {user['session_limit']} seconds ({user['session_limit']//3600}h {(user['session_limit']%3600)//60}m)
  MAC Addresses: {user['mac_addresses'] or 'None'}
  Status: {'Active' if user['is_active'] else 'Inactive'}
""")

            logger.info("="*80)
        else:
            logger.info("No users found in database")

        return True

    except Exception as e:
        logger.error(f"Error displaying created data: {e}")
        return False


def main():
    """Main function to seed the database."""
    global logger

    # Setup logging
    logger = setup_logging(log_level='INFO')

    logger.info("Starting database seeding...")

    # Check database connectivity
    if not db.health_check():
        logger.error("Database connection failed")
        sys.exit(1)

    # Create sample data
    if not create_sample_users():
        logger.error("Failed to create sample users")
        sys.exit(1)

    # Display created data
    if not display_created_data():
        logger.error("Failed to display created data")
        sys.exit(1)

    logger.info("Database seeding completed successfully!")
    logger.info("\nYou can now test the VPN access server with these sample users:")
    logger.info("- Username: john.doe, MAC: 00:11:22:33:44:55 or 00:11:22:33:44:56")
    logger.info("- Username: jane.smith, MAC: aa:bb:cc:dd:ee:ff")
    logger.info("- Username: admin, MAC: 12:34:56:78:9a:bc, 12:34:56:78:9a:bd, or 12:34:56:78:9a:be")


if __name__ == "__main__":
    main()