n3ro-cli/main.py

import base64
import csv
import sys
import os
import urllib.request
from urllib.parse import parse_qs


class N3ROCLI:

    def __init__(self):
        self.config_path = os.path.expanduser('~/.config/n3ro-cli')
        self.ssr_command = f"python {self.config_path}/shadowsocksr/shadowsocks/local.py"
        self.url_path = os.path.join(self.config_path, 'url')
        self.cache_path = os.path.join(self.config_path, 'cache')
        self.subscribe_url = self.read_url()
        os.makedirs(self.config_path, exist_ok=True)

    def read_url(self):
        path_ok = os.path.exists(self.url_path)
        if path_ok:
            with open(self.url_path) as u:
                url = u.read()
            return url
        else:
            return None

    def decode_str(self, s, to_str=True):
        paddings = 4 - len(s) % 4
        padded_s = f"{s}{'='*paddings}"
        decoded = base64.urlsafe_b64decode(padded_s)
        if to_str:
            decoded = str(decoded, encoding='utf-8')
        return decoded

    def sync(self):
        if not self.subscribe_url:
            print("add a ssr subscribe url in ~/.config/n3ro-cli/url")
        else:
            with urllib.request.urlopen(self.subscribe_url) as resp:
                data = str(resp.read(), encoding='utf-8')
            ssr_urls = self.decode_str(data).strip().split('\n')
            ssr_paths = (self.decode_str(u[6:]) for u in ssr_urls)
            with open(self.cache_path, 'w') as c:
                fields = ['id', 'info', 'args']
                server_writer = csv.DictWriter(c, fieldnames=fields)
                server_writer.writeheader()
                for n, p in enumerate(ssr_paths):
                    path, params = p.split('/?')
                    d = self.decode_str
                    host, port, proto, enc, obfs, passwd = path.split(':')
                    raw_passwd = d(passwd)
                    parsed_params = parse_qs(params)
                    obfs_param = parsed_params.get('obfsparam', '')
                    if obfs_param:
                        obfs_param = d(obfs_param)
                    proto_param = parsed_params.get('proto_param', '')
                    if proto_param:
                        proto_param = d(proto_param)
                    name = d(parsed_params.get('remarks', ''),
                             to_str=False).decode()
                    group = d(parsed_params.get('group', ''),
                              to_str=False).decode()
                    args = f"-s {host} -p {port} -k {raw_passwd} -m {enc} -O {proto} -o {obfs} -b 127.0.0.1 -l 1080"
                    if obfs_param:
                        args = f"{args} -g {obfs_param}"
                    if proto_param:
                        args = f"{args} -G {proto_param}"
                    server_writer.writerow({'id': n,
                                            'args': args,
                                            'info': f"{name} {group}"})

    def conn(self, node_id):
        with open(self.cache_path) as c:
            reader = csv.DictReader(c)
            ssr_args = None
            for row in reader:
                if row['id'] == str(node_id):
                    ssr_args = row['args']
        os.system(f"{self.ssr_command} {ssr_args}")

    def ls(self):
        with open(self.cache_path) as c:
            reader = csv.DictReader(c)
            for row in reader:
                s = f"{row['id']}:{row['info']}  "
                sys.stdout.write(s)


cli = N3ROCLI()
ls = cli.ls
l = cli.ls
conn = cli.conn
c = cli.conn
sync = cli.sync
s = cli.sync