App/Mikrotik/utils/api.py

from librouteros import connect
from librouteros.query import Key
import logging
from typing import Optional
from datetime import datetime
from ..schema import *
from ..schema import UserListData
from .helperfx import PhoneValidator
from .Config import RouterConfig
from time import sleep


class MikrotikAPI:
    """Class to handle MikroTik router API operations, inherits VPNManager."""

    def __init__(
        self,
        config: RouterConfig,
    ):
        """Initialize with router configuration and VPN configuration."""
        self.config = config
        self.api = None
        self.logger = self._setup_logging()
        self.phone_validator = PhoneValidator()

    def connect(self) -> bool:
        """Establish connection to VPN and MikroTik router."""

        try:
            self.api = connect(
                username=self.config.username,
                password=self.config.password,
                host=self.config.host,
                port=8728,
            )
            self.logger.info("Successfully connected to MikroTik router.")
            return True
        except Exception as e:
            self.logger.error(f"Failed to connect to MikroTik router: {e}")
            return False

    def close(self):
        """Close the router connection and disconnect VPN."""
        if self.api:
            self.api.close()
            self.logger.info("Router connection closed.")

    def add_hotspot_user(
        self, phone: str, password: str, profile: str = "default"
    ) -> BaseResponse:
        """Add a hotspot user to the router"""
        if not self.api:
            return ResponseBuilder.error(
                BaseResponse, "Registration failed", "Not connected to router"
            )

        # Validate phone number
        is_valid, formatted_phone = self.phone_validator.validate_and_format(phone)
        if not is_valid:
            return ResponseBuilder.error(
                BaseResponse,
                "Registration failed",
                f"Invalid phone number format: {phone}",
            )

        users = self.api.path("ip", "hotspot", "user")
        existing_users = users.select(Key("name"))

        # Check if user already exists
        if any(str(user["name"]) == formatted_phone for user in existing_users):
            return ResponseBuilder.error(
                BaseResponse,
                "Registration failed",
                f"User {formatted_phone} already exists",
            )

        try:
            # Add user to router
            users.add(
                name=formatted_phone,
                password=password,
                profile=profile,
                disabled=True,
            )
            return ResponseBuilder.success(BaseResponse, "User registered successfully")

        except Exception as e:
            self.logger.error(f"Failed to add user {formatted_phone}: {e}")
            return ResponseBuilder.error(BaseResponse, "Registration failed", str(e))

    def login_hotspot_user(
        self, phone: str, password: str, mac_address: str, ip_address: str
    ) -> LoginResponse:
        """Login a hotspot user"""
        if not self.api:
            return ResponseBuilder.error(
                LoginResponse, "Login failed", "Not connected to router"
            )

        is_valid, formatted_phone = self.phone_validator.validate_and_format(phone)
        if not is_valid:
            return ResponseBuilder.error(
                LoginResponse, "Login failed", f"Invalid phone number format: {phone}"
            )

        try:
            users = self.api.path("ip", "hotspot", "user")
            active = self.api.path("ip", "hotspot", "active")

            # Verify if user exists and password is correct
            existing_users = users.select(Key("name"), Key("password"), Key("profile"))
            user_data = next(
                (
                    user
                    for user in existing_users
                    if str(user["name"]) == formatted_phone
                    and user["password"] == password
                ),
                None,
            )

            if not user_data:
                return ResponseBuilder.error(
                    LoginResponse, "Login failed", "Invalid credentials"
                )

            # Check if already logged in
            active_users = active.select(Key("user"))
            if any(str(user["user"]) == formatted_phone for user in active_users):
                return ResponseBuilder.error(
                    LoginResponse,
                    "Login failed",
                    f"User {formatted_phone} is already logged in",
                )

            # Perform the login
            self.api.path("ip", "hotspot", "host").call(
                "login",
                user=formatted_phone,
                password=password,
                mac_address=mac_address,
                ip_address=ip_address,
            )
            login_time = datetime.now()
            return LoginResponse(
                code=200, message="Login successful", login_time=login_time
            )

        except Exception as e:
            self.logger.error(f"Login failed for {formatted_phone}: {e}")
            return ResponseBuilder.error(LoginResponse, "Login failed", str(e))

    def logout_hotspot_user(self, phone: str) -> LogoutResponse:
        """Logout a hotspot user"""
        if not self.api:
            return ResponseBuilder.error(
                LogoutResponse, "Logout failed", "Not connected to router"
            )

        is_valid, formatted_phone = self.phone_validator.validate_and_format(phone)
        if not is_valid:
            return ResponseBuilder.error(
                LogoutResponse, "Logout failed", f"Invalid phone number format: {phone}"
            )

        active = self.api.path("ip", "hotspot", "active")

        try:
            active_users = active.select(Key("user"), Key(".id"))
            user_session = next(
                (
                    session
                    for session in active_users
                    if str(session["user"]) == formatted_phone
                ),
                None,
            )

            if not user_session:
                return ResponseBuilder.error(
                    LogoutResponse,
                    "Logout failed",
                    f"User {formatted_phone} is not logged in",
                )

            # Logout the user
            active.remove(user_session[".id"])
            logout_time = datetime.now()
            return LogoutResponse(
                code=200, message="Logout successful", logout_time=logout_time
            )

        except Exception as e:
            self.logger.error(f"Logout failed for user {formatted_phone}: {e}")
            return ResponseBuilder.error(LogoutResponse, "Logout failed", str(e))

    def get_active_users(self) -> ActiveUsersResponse:
        """Get list of active hotspot users"""
        if not self.api:
            return ResponseBuilder.error(
                ActiveUsersResponse,
                "Failed to get active users",
                "Not connected to router",
            )

        try:
            active = self.api.path("ip", "hotspot", "active")
            active_users = active.select(
                Key("user"),
                Key("uptime"),
                Key("mac-address"),
                Key("address"),
                Key("bytes-in"),
                Key("bytes-out"),
                # Key("profile"),
            )

            users_list = [
                UserSessionData(
                    phone=str(user["user"]),
                    uptime=user["uptime"],
                    mac_address=user["mac-address"],
                    ip_address=user["address"],
                    bytes_in=int(user["bytes-in"]),
                    bytes_out=int(user["bytes-out"]),
                    status="active",
                    profile=str("2mbps_profile"),
                )
                for user in active_users
            ]
            print(users_list)
            return ActiveUsersResponse(
                code=200,
                message="Active users retrieved successfully",
                active_users=users_list,
            )

        except Exception as e:
            self.logger.error(f"Failed to get active users: {e}")
            return ResponseBuilder.error(
                ActiveUsersResponse, "Failed to get active users", str(e)
            )

    def get_user_stats(self, phone: Optional[str] = None) -> UserStatsResponse:
        """Get statistics for all users or a specific user"""
        if not self.api:
            return ResponseBuilder.error(
                UserStatsResponse,
                "Failed to get user statistics",
                "Not connected to router",
            )

        try:
            users = self.api.path("ip", "hotspot", "user")
            all_users = users.select(
                Key("name"),
                Key("profile"),
                Key("disabled"),
                Key("uptime"),
                Key("bytes-in"),
                Key("bytes-out"),
            )

            user_stats = [
                UserStats(
                    phone=str(user["name"]),
                    profile=user.get("profile", "default"),
                    status=(
                        "disabled"
                        if user.get("disabled", "false") == "true"
                        else "enabled"
                    ),
                    uptime=user.get("uptime", "0s"),
                    bytes_in=int(user.get("bytes-in", 0)),
                    bytes_out=int(user.get("bytes-out", 0)),
                )
                for user in all_users
                if not phone or str(user["name"]) == phone
            ]

            if phone and not user_stats:
                return ResponseBuilder.error(
                    UserStatsResponse, "User not found", f"User {phone} does not exist"
                )

            return UserStatsResponse(
                code=200,
                message="User statistics retrieved successfully",
                user_stats=user_stats,
            )

        except Exception as e:
            self.logger.error(f"Failed to get user statistics: {e}")
            return ResponseBuilder.error(
                UserStatsResponse, "Failed to get user statistics", str(e)
            )

    # Existing initialization, connect, and other methods...

    def set_user_status(self, phone: str, disabled: bool) -> UserStatusResponse:
        """Enable or disable a hotspot user."""
        if not self.api:
            return ResponseBuilder.error(
                UserStatusResponse, "Status update failed", "Not connected to router"
            )

        # Validate phone number
        is_valid, formatted_phone = self.phone_validator.validate_and_format(phone)
        if not is_valid:
            return ResponseBuilder.error(
                UserStatusResponse,
                "Status update failed",
                f"Invalid phone number format: {phone}",
            )
        users = self.api.path("ip", "hotspot", "user")

        try:
            # Find the user
            existing_users = users.select(
                Key("name"),
                Key(".id"),
                Key("mbonea"),
            )
            for user in existing_users:
                if str(user["name"]) == formatted_phone:
                    user_data = user
                    break

            if not user_data:
                return ResponseBuilder.error(
                    UserStatusResponse,
                    "Status update failed",
                    f"User {formatted_phone} does not exist",
                )

            user_id = user_data[".id"]
            status_value = "true" if disabled else "false"
            users.update(**{".id": user_id, "disabled": status_value})
            status = "disabled" if disabled else "enabled"
            self.logger.info(f"User {formatted_phone} successfully {status}")

            # If disabling the user, log them out of any active session
            if disabled:
                self.logout_hotspot_user(formatted_phone)

            return UserStatusResponse(
                code=200, message=f"User successfully {status}", user_status=status
            )

        except Exception as e:
            print(f"Failed to update status for user {formatted_phone}: {e}")
            self.logger.error(
                f"Failed to update status for user {formatted_phone}: {e}"
            )
            return ResponseBuilder.error(
                UserStatusResponse, "Status update failed", str(e)
            )

    def activate_user(self, phone: str) -> UserStatusResponse:
        """Activate a user (set disabled to false)."""
        return self.set_user_status(phone, disabled=False)

    def deactivate_user(self, phone: str) -> UserStatusResponse:
        """Deactivate a user (set disabled to true and log them out)."""
        return self.set_user_status(phone, disabled=True)

    def get_users_list(self) -> UsersListResponse:
        """Get list of all hotspot users (both active and inactive)"""
        if not self.api:
            return UsersListResponse(
                success=False,
                message="Failed to get users list",
                error="Not connected to router",
            )

        try:
            users = self.api.path("ip", "hotspot", "user")
            active = self.api.path("ip", "hotspot", "active")

            # Retrieve all users
            all_users = users.select(
                Key("name"),
                Key("profile"),
                Key("disabled"),
                Key("comment"),
                Key("limit-bytes-total"),
                Key("bytes-in"),
                Key("bytes-out"),
            )

            # Retrieve active users to check statuses
            active_users = active.select(Key("user"))
            active_phones = [str(user.get("user", "")) for user in active_users]

            # Format the response
            users_list = []
            for user in all_users:
                phone_number = str(user["name"])
                bytes_in = int(user.get("bytes-in", 0))
                bytes_out = int(user.get("bytes-out", 0))
                data_limit = (
                    int(user.get("limit-bytes-total", 0))
                    if user.get("limit-bytes-total")
                    else None
                )

                user_data = UserListData(
                    phoneNumber=phone_number,
                    profile=user.get("profile", "default"),
                    status="active" if phone_number in active_phones else "inactive",
                    disabled=user.get("disabled", "false") == "true",
                    comment=user.get("comment", ""),
                    data_limit=data_limit,
                )
                users_list.append(user_data)

            return UsersListResponse(
                code=200,
                message=f"Retrieved {len(users_list)} users",
                data=users,
            )

        except Exception as e:
            return UsersListResponse(
                success=False, message="Failed to get users list", error=str(e)
            )