octo-funnel/octoprint_tailscale_funnel/tailscale_interface.py

# coding=utf-8
from __future__ import absolute_import

import subprocess
import json
import re


class TailscaleError(Exception):
    """Base exception for Tailscale interface errors"""
    pass


class TailscaleNotInstalledError(TailscaleError):
    """Exception raised when Tailscale is not installed"""
    pass


class TailscaleNotRunningError(TailscaleError):
    """Exception raised when Tailscale is not running"""
    pass


class FunnelNotEnabledError(TailscaleError):
    """Exception raised when Funnel is not enabled"""
    pass


class TailscaleInterface:
    def __init__(self, logger=None):
        self._logger = logger

    def _run_command(self, command):
        """
        Run a command and return the result
        """
        try:
            result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
            if result.returncode == 0:
                return {"success": True, "output": result.stdout.strip(), "error": None}
            else:
                return {"success": False, "output": None, "error": result.stderr.strip()}
        except subprocess.TimeoutExpired:
            if self._logger:
                self._logger.error("Command timed out: '{}'".format(command))
            return {"success": False, "output": None, "error": "Command timed out"}
        except Exception as e:
            if self._logger:
                self._logger.error("Error running command '{}': {}".format(command, str(e)))
            return {"success": False, "output": None, "error": str(e)}

    def is_tailscale_installed(self):
        """
        Check if tailscale is installed
        """
        result = self._run_command("which tailscale")
        return result["success"] and result["output"] != ""

    def is_tailscale_running(self):
        """
        Check if tailscale is running
        """
        result = self._run_command("tailscale status --json")
        return result["success"]

    def get_tailscale_status(self):
        """
        Get tailscale status as JSON
        """
        if not self.is_tailscale_installed():
            raise TailscaleNotInstalledError("Tailscale is not installed")
            
        result = self._run_command("tailscale status --json")
        if result["success"]:
            try:
                return json.loads(result["output"])
            except json.JSONDecodeError as e:
                if self._logger:
                    self._logger.error("Error parsing tailscale status JSON: {}".format(str(e)))
                raise TailscaleError("Error parsing tailscale status: {}".format(str(e)))
        else:
            raise TailscaleNotRunningError("Tailscale is not running: {}".format(result["error"]))

    def is_funnel_enabled(self):
        """
        Check if tailscale funnel is enabled
        """
        if not self.is_tailscale_installed():
            raise TailscaleNotInstalledError("Tailscale is not installed")
            
        result = self._run_command("tailscale funnel status --json")
        if result["success"]:
            try:
                status = json.loads(result["output"])
                # Check if any funnel is active
                for _, ports in status.get("AllowFunnel", {}).items():
                    if any(ports.values()):
                        return True
                return False
            except json.JSONDecodeError as e:
                if self._logger:
                    self._logger.error("Error parsing funnel status JSON: {}".format(str(e)))
                raise TailscaleError("Error parsing funnel status: {}".format(str(e)))
        else:
            # If the command fails, it might be because Funnel is not enabled
            # Check if it's a specific error we can handle
            if "funnel" in result["error"].lower() and "not enabled" in result["error"].lower():
                return False
            raise TailscaleError("Error checking funnel status: {}".format(result["error"]))

    def enable_funnel(self, port):
        """
        Enable tailscale funnel for the specified port
        """
        if not self.is_tailscale_installed():
            raise TailscaleNotInstalledError("Tailscale is not installed")
            
        command = "tailscale funnel {} on".format(port)
        result = self._run_command(command)
        return result["success"]

    def disable_funnel(self, port):
        """
        Disable tailscale funnel for the specified port
        """
        if not self.is_tailscale_installed():
            raise TailscaleNotInstalledError("Tailscale is not installed")
            
        command = "tailscale funnel {} off".format(port)
        result = self._run_command(command)
        return result["success"]

    def get_public_url(self):
        """
        Get the public URL for accessing the service via funnel
        """
        if not self.is_tailscale_installed():
            raise TailscaleNotInstalledError("Tailscale is not installed")
            
        result = self._run_command("tailscale funnel status --json")
        if result["success"]:
            try:
                status = json.loads(result["output"])
                # Extract the public URL
                # This is a simplified approach - in practice, you might need to parse more details
                funnel_status = status.get("AllowFunnel", {})
                if funnel_status:
                    # Get the first available funnel URL
                    for key, ports in funnel_status.items():
                        if any(ports.values()):
                            # Extract the domain from the key
                            match = re.search(r'https://([a-zA-Z0-9\-\.]+)', key)
                            if match:
                                return "https://{}".format(match.group(1))
                return None
            except (json.JSONDecodeError, KeyError) as e:
                if self._logger:
                    self._logger.error("Error parsing funnel status for URL: {}".format(str(e)))
                raise TailscaleError("Error parsing funnel status for URL: {}".format(str(e)))
        else:
            raise TailscaleError("Error getting funnel status: {}".format(result["error"]))