KTVManager_UI/test_npm_update.py

import requests
import json
import mysql.connector
import re
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Assuming config.py is in the same directory or accessible via PYTHONPATH
from config import DevelopmentConfig as app_config

class NginxProxyManager:
    def __init__(self, host, email, password):
        self.host = host
        self.email = email
        self.password = password
        self.token = None

    def login(self):
        url = f"{self.host}/api/tokens"
        payload = {
            "identity": self.email,
            "secret": self.password
        }
        headers = {
            "Content-Type": "application/json"
        }
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        if response.status_code == 200:
            self.token = response.json()["token"]
            print("Login successful.")
        else:
            print(f"Failed to login: {response.text}")
            exit(1)

    def get_proxy_host(self, host_id):
        if not self.token:
            self.login()

        url = f"{self.host}/api/nginx/proxy-hosts/{host_id}"
        headers = {
            "Authorization": f"Bearer {self.token}"
        }
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.json()
        else:
            print(f"Failed to get proxy host {host_id}: {response.text}")
            return None

    def update_proxy_host_config(self, host_id, config):
        if not self.token:
            self.login()

        url = f"{self.host}/api/nginx/proxy-hosts/{host_id}"
        
        original_host_data = self.get_proxy_host(host_id)
        if not original_host_data:
            return

        # Construct a new payload with only the allowed fields for an update
        update_payload = {
            "domain_names": original_host_data.get("domain_names", []),
            "forward_scheme": original_host_data.get("forward_scheme", "http"),
            "forward_host": original_host_data.get("forward_host"),
            "forward_port": original_host_data.get("forward_port"),
            "access_list_id": original_host_data.get("access_list_id", 0),
            "certificate_id": original_host_data.get("certificate_id", 0),
            "ssl_forced": original_host_data.get("ssl_forced", False),
            "hsts_enabled": original_host_data.get("hsts_enabled", False),
            "hsts_subdomains": original_host_data.get("hsts_subdomains", False),
            "http2_support": original_host_data.get("http2_support", False),
            "block_exploits": original_host_data.get("block_exploits", False),
            "caching_enabled": original_host_data.get("caching_enabled", False),
            "allow_websocket_upgrade": original_host_data.get("allow_websocket_upgrade", False),
            "advanced_config": config, # The updated advanced config
            "meta": original_host_data.get("meta", {}),
            "locations": original_host_data.get("locations", []),
        }
        
        headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }

        response = requests.put(url, headers=headers, data=json.dumps(update_payload))
        if response.status_code == 200:
            print(f"Successfully updated proxy host {host_id}")
        else:
            print(f"Failed to update proxy host {host_id}: {response.text}")


def get_streams_from_db(db_host, db_user, db_pass, db_name, db_port):
    try:
        conn = mysql.connector.connect(
            host=db_host,
            user=db_user,
            password=db_pass,
            database=db_name,
            port=db_port
        )
        cursor = conn.cursor(dictionary=True)
        cursor.execute("""
            SELECT DISTINCT
              SUBSTRING_INDEX(stream, ' ', 1) AS streamName,
              streamURL
            FROM userAccounts
        """)
        streams = cursor.fetchall()
        cursor.close()
        conn.close()
        return streams
    except mysql.connector.Error as err:
        print(f"Error connecting to database: {err}")
        return []

def update_config_with_streams(config, streams):
    # Get all stream names from the database
    db_stream_names = {stream['streamName'] for stream in streams}

    # Find all location blocks in the config
    location_blocks = re.findall(r'location ~ \^/(\w+)\(\.\*\)\$ \{[^}]+\}', config)

    # Remove location blocks that are not in the database
    for stream_name in location_blocks:
        if stream_name not in db_stream_names:
            print(f"Removing location block for stream: {stream_name}")
            pattern = re.compile(f'location ~ \\^/{re.escape(stream_name)}\\(\\.\\*\\)\\$ {{[^}}]+}}\\s*', re.DOTALL)
            config = pattern.sub('', config)

    # Update existing stream URLs
    for stream in streams:
        stream_name = stream['streamName']
        stream_url = stream['streamURL']
        if stream_url: # Ensure there is a URL to update to
            # Use a more specific regex to avoid replacing parts of other URLs
            pattern = re.compile(f'(location ~ \\^/{re.escape(stream_name)}\\(\\.\\*\\)\\$ {{\\s*return 302 )([^;]+)(;\\s*}})')
            config = pattern.sub(f'\\1{stream_url}/$1$is_args$args\\3', config)
            
    return config

def main():
    npm = NginxProxyManager(app_config.NPM_HOST, app_config.NPM_EMAIL, app_config.NPM_PASSWORD)
    npm.login()
    host = npm.get_proxy_host(9)
    if host:
        current_config = host.get('advanced_config', '')
        print("Current Config:")
        print(current_config)
        
        streams = get_streams_from_db(app_config.DBHOST, app_config.DBUSER, app_config.DBPASS, app_config.DATABASE, app_config.DBPORT)
        if streams:
            new_config = update_config_with_streams(current_config, streams)
            print("\nNew Config:")
            print(new_config)
            
            # Uncomment the following line to apply the changes
            npm.update_proxy_host_config(9, new_config)
            print("\nTo apply the changes, uncomment the last line in the main function.")

if __name__ == "__main__":
    main()