omnigen2/utils/vpn_utils.py

import random

import json
import yaml

import requests

class VPNManager:
    def __init__(self, config_path: str = '/etc/mihomo/config.yaml'):
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
        self.external_controller = config['external-controller']
        self.external_controller = self.external_controller.replace('0.0.0.0', '127.0.0.1')
        self.secret = config['secret']

        self.headers = {"Authorization": f"Bearer {self.secret}"}

        self.unavailable_nodes = set()
    
    @property
    def current_node(self):
        url = f"http://{self.external_controller}/group/Proxy"
        r = requests.request("GET", url, headers=self.headers)
        return r.json()['now']
    
    @property
    def available_nodes(self):
        return list(self.get_available_vpn_nodes() - self.unavailable_nodes)

    def switch_vpn_node(self, node_name):
        url = f"http://{self.external_controller}/proxies/Proxy"

        payload = json.dumps({
            "name": node_name
        })
        headers = self.headers.copy()
        headers.update({'Content-Type': 'application/json'})
        r = requests.request("PUT", url, headers=headers, data=payload)
        if r.status_code != 204:
            raise Warning(f"Failed to switch to {node_name}")
        return r.status_code == 204

    def get_random_available_vpn_node(self):
        return random.choice(self.available_nodes)
    
    def random_switch_vpn_node(self):
        node_name = self.get_random_available_vpn_node()
        print(f"Switching to {node_name}")
        self.switch_vpn_node(node_name)
        # self.current_node = node_name
        return node_name
    
    def get_vpn_nodes(self):
        url = f"http://{self.external_controller}/group/Proxy"
        delay_res = requests.get(url, headers=self.headers)
        return delay_res.json()['all']
    
    def get_available_vpn_nodes(self):
        url = f"http://{self.external_controller}/group/Proxy/delay?timeout=5000&url=http://www.gstatic.com/generate_204"
        delay_res = requests.get(url, headers=self.headers)
        return set(delay_res.json().keys())
    
    def get_current_vpn_node_ip(self):
        url = "http://ifconfig.me"
        r = requests.request("GET", url)
        return r.text

    def add_unavailable_node(self, node_name):
        self.unavailable_nodes.add(node_name)