vpn-btcpay-provisioner/scripts/subscription_checker.py

import logging
import subprocess
import os
import tempfile
from pathlib import Path
from datetime import datetime, timedelta
from utils.db.operations import DatabaseManager
from utils.db.models import SubscriptionStatus
from app.handlers.payment_handler import BTCPayHandler

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Path setup
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
CLEANUP_PLAYBOOK = PROJECT_ROOT / 'ansible' / 'playbooks' / 'vpn_cleanup.yml'
INVENTORY_FILE = PROJECT_ROOT / 'inventory.ini'

# Notification thresholds configuration
NOTIFICATION_THRESHOLDS = {
    'minimum_duration': timedelta(hours=1),  # Minimum subscription duration
    'short_term': {
        'max_duration': timedelta(days=1),
        'warning_fraction': 0.5,  # Warn when 50% of time remains
        'grace_fraction': 0.1     # Grace period of 10% of subscription length
    },
    'medium_term': {
        'max_duration': timedelta(days=7),
        'warning_fraction': 0.25,  # Warn when 25% of time remains
        'grace_hours': 12         # Fixed 12-hour grace period
    },
    'long_term': {
        'warning_days': 7,        # 7 days warning for longer subscriptions
        'grace_days': 2          # 2 days grace period
    }
}

def run_cleanup_playbook(subscription_id):
    """Run the VPN cleanup playbook"""
    logger.info(f"Running cleanup playbook for subscription {subscription_id}")
    
    vault_pass = os.getenv('ANSIBLE_VAULT_PASSWORD', '')
    if not vault_pass:
        raise Exception("Vault password not found in environment variables")
    
    with tempfile.NamedTemporaryFile(mode='w', delete=False) as vault_pass_file:
        vault_pass_file.write(vault_pass)
        vault_pass_file.flush()
        
        cmd = [
            'ansible-playbook',
            str(CLEANUP_PLAYBOOK),
            '-i', str(INVENTORY_FILE),
            '-e', f'subscription_id={subscription_id}',
            '--vault-password-file', vault_pass_file.name,
            '-vvv'
        ]
        
        logger.debug(f"Running command: {' '.join(cmd)}")
        
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            logger.error(f"Cleanup playbook failed: {result.stderr}")
        else:
            logger.debug(f"Cleanup playbook output: {result.stdout}")
    
    os.unlink(vault_pass_file.name)
    return result

def calculate_notification_times(start_time, end_time):
    """
    Calculate warning and grace period times based on subscription duration
    
    Returns:
        tuple: (warning_time, grace_end_time)
    """
    duration = end_time - start_time
    
    # Handle extremely short subscriptions
    if duration < NOTIFICATION_THRESHOLDS['minimum_duration']:
        warning_time = start_time  # Immediate warning
        grace_end_time = end_time  # No grace period
        return warning_time, grace_end_time
    
    # Short-term subscriptions (less than 1 day)
    if duration <= NOTIFICATION_THRESHOLDS['short_term']['max_duration']:
        warning_delta = duration * NOTIFICATION_THRESHOLDS['short_term']['warning_fraction']
        grace_delta = duration * NOTIFICATION_THRESHOLDS['short_term']['grace_fraction']
        warning_time = end_time - warning_delta
        grace_end_time = end_time + grace_delta
        
    # Medium-term subscriptions (1-7 days)
    elif duration <= NOTIFICATION_THRESHOLDS['medium_term']['max_duration']:
        warning_delta = duration * NOTIFICATION_THRESHOLDS['medium_term']['warning_fraction']
        grace_hours = NOTIFICATION_THRESHOLDS['medium_term']['grace_hours']
        warning_time = end_time - warning_delta
        grace_end_time = end_time + timedelta(hours=grace_hours)
        
    # Long-term subscriptions (> 7 days)
    else:
        warning_days = NOTIFICATION_THRESHOLDS['long_term']['warning_days']
        grace_days = NOTIFICATION_THRESHOLDS['long_term']['grace_days']
        warning_time = end_time - timedelta(days=warning_days)
        grace_end_time = end_time + timedelta(days=grace_days)
    
    return warning_time, grace_end_time

def get_notification_message(subscription, remaining_time):
    """Generate appropriate notification message based on subscription duration"""
    if remaining_time < timedelta(hours=1):
        return f"Your VPN subscription expires in {int(remaining_time.total_seconds() / 60)} minutes!"
    elif remaining_time < timedelta(days=1):
        return f"Your VPN subscription expires in {int(remaining_time.total_seconds() / 3600)} hours!"
    else:
        return f"Your VPN subscription expires in {remaining_time.days} days!"

def notify_user(subscription, message):
    """Send notification to user about subscription status"""
    try:
        if not subscription.user or not subscription.user.email:
            logger.error(f"No email found for subscription {subscription.id}")
            return False
        
        btcpay_handler = BTCPayHandler()
        email_sent = btcpay_handler.send_confirmation_email(
            subscription.user.email,
            f"""
            VPN Subscription Update
            
            {message}
            
            If you wish to continue using the VPN service, please renew your subscription.
            """
        )
        
        if email_sent:
            logger.info(f"Sent notification to {subscription.user.email}: {message}")
        else:
            logger.warning(f"Failed to send notification to {subscription.user.email}")
            
        return email_sent
        
    except Exception as e:
        logger.error(f"Error sending notification: {str(e)}")
        return False

def check_subscriptions():
    """Check subscription status and clean up expired ones"""
    logger.info("Starting subscription check")
    
    try:
        active_subscriptions = DatabaseManager.get_active_subscriptions()
        logger.info(f"Checking {len(active_subscriptions)} active subscriptions")
        
        now = datetime.utcnow()
        
        for subscription in active_subscriptions:
            try:
                logger.debug(f"Processing subscription {subscription.id}")
                
                warning_time, grace_end_time = calculate_notification_times(
                    subscription.start_time,
                    subscription.expiry_time
                )
                
                # Calculate remaining time
                remaining_time = subscription.expiry_time - now
                logger.debug(f"Subscription {subscription.id} has {remaining_time} remaining")
                
                # Handle warnings
                if now >= warning_time and not subscription.warning_sent:
                    message = get_notification_message(subscription, remaining_time)
                    logger.info(f"Sending notification for subscription {subscription.id}: {message}")
                    
                    if notify_user(subscription, message):
                        DatabaseManager.update_warning_sent(subscription.id)
                
                # Handle expiration
                if now >= grace_end_time:
                    logger.info(f"Processing expiration for subscription {subscription.id}")
                    
                    try:
                        result = run_cleanup_playbook(subscription.invoice_id)
                        
                        if result.returncode == 0:
                            DatabaseManager.expire_subscription(subscription.id)
                            logger.info(f"Successfully cleaned up subscription {subscription.id}")
                            
                            # Send final notification
                            notify_user(subscription, "Your VPN subscription has expired and been deactivated.")
                        else:
                            logger.error(f"Cleanup failed: {result.stderr}")
                            
                    except Exception as e:
                        logger.error(f"Error during cleanup: {str(e)}")
                        logger.error(traceback.format_exc())
                        
            except Exception as e:
                logger.error(f"Error processing subscription {subscription.id}: {str(e)}")
                logger.error(traceback.format_exc())
                continue
        
        logger.info("Subscription check completed")
        
    except Exception as e:
        logger.error(f"Subscription checker failed: {str(e)}")
        logger.error(traceback.format_exc())
        raise

if __name__ == '__main__':
    try:
        check_subscriptions()
    except Exception as e:
        logger.error(f"Subscription checker failed: {str(e)}")
        raise