#!/usr/bin/env python3 from __future__ import annotations import typing import sys import os import json from pathlib import Path from ipaddress import IPv4Network, IPv6Network, AddressValueError import ray_pb2 def load_networks( dat_files: typing.List[Path], countries: typing.List[str], inverse: bool = False ) -> typing.List[IPv4Network | IPv6Network]: networks = [] country_codes = [code.lower() for code in countries] for file in dat_files: file_path = file.resolve() if not file_path.exists() or not file_path.is_file(): continue data = ray_pb2.GeoIPList() with open(file_path, 'rb') as fp: data.ParseFromString(fp.read()) for entry in data.entry: if inverse: if entry.country_code.lower() in country_codes: continue else: if entry.country_code.lower() not in country_codes: continue for cidr in entry.cidr: try: networks.append(IPv4Network((cidr.ip, cidr.prefix))) except AddressValueError: pass try: networks.append(IPv6Network((cidr.ip, cidr.prefix))) except AddressValueError: pass return networks def load_domains( dat_files: typing.List[Path], countries: typing.List[str], inverse: bool = False, formatter: Path | None = None, domain_types: typing.List[int] | None = None, ) -> typing.List[str]: domains = [] country_codes = [code.lower() for code in countries] format_dict = None if formatter: format_file_path = formatter.resolve() if format_file_path.is_file(): with format_file_path.open(encoding='utf-8') as fp: format_dict = json.load(fp) for file in dat_files: file_path = file.resolve() if not file_path.exists() or not file_path.is_file(): continue data = ray_pb2.GeoSiteList() with open(file_path, 'rb') as fp: data.ParseFromString(fp.read()) for entry in data.entry: if inverse: if entry.country_code.lower() in country_codes: continue else: if entry.country_code.lower() not in country_codes: continue for domain in entry.domain: value = None prefix = '' suffix = '' if format_dict: prefix = format_dict.get('prefix').get(str(domain.type)) suffix = format_dict.get('suffix').get(str(domain.type)) if domain_types is None or domain.type in domain_types: value = domain.value if value is not None: domains.append(f'{prefix}{domain.value}{suffix}') return domains def show_data( data_type: str, dat_files: typing.List[Path], countries: typing.List[str], inverse: bool = False, ) -> None: country_codes = [code.lower() for code in countries] for file in dat_files: file_path = file.resolve() if not file_path.exists() or not file_path.is_file(): continue data = None if data_type == 'geoip': data = ray_pb2.GeoIPList() list_name = 'cidr' if data_type == 'geosite': data = ray_pb2.GeoSiteList() list_name = 'domain' if not data: return with open(file_path, 'rb') as fp: data.ParseFromString(fp.read()) for entry in data.entry: if len(country_codes) > 0: if inverse: if entry.country_code.lower() in country_codes: continue else: if entry.country_code.lower() not in country_codes: continue print(f"{entry.country_code}: {len(getattr(entry, list_name, None))}") def main(args): input_files = getattr(args, 'input', []) output = getattr(args, 'output', None) if getattr(args, 'show', False): show_data( getattr(args, 'data_type', None), input_files, getattr(args, 'country', []), getattr(args, 'inverse', False), ) sys.exit(0) data_list = [] if getattr(args, 'data_type', None) == 'geoip': data_list = load_networks( input_files, getattr(args, 'country', []), getattr(args, 'inverse', False), ) if getattr(args, 'data_type', None) == 'geosite': data_list = load_domains( input_files, getattr(args, 'country', []), getattr(args, 'inverse', None), getattr(args, 'formatter', None), getattr(args, 'domain_types', None), ) if len(data_list) == 0: print('No legal data.') sys.exit(1) output_lines = [ f"{getattr(args,'prefix','')}{str(entry)}{getattr(args,'suffix','')}" for entry in data_list ] if output is None: for line in output_lines: print(line, end='') else: output: Path output_file = output.resolve() if not output_file.parent.exists(): os.makedirs(output_file.parent) with output_file.open('w', encoding='utf-8') as wfp: wfp.writelines(output_lines) print(f'Output: {len(output_lines)}') if __name__ == '__main__': import argparse parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( 'input', nargs='+', type=Path, default=[], help='input ray data files (separated by space)', ) parser.add_argument( '-c', '--country', type=str, action='append', default=[], help='country codes (use multiple times)', ) parser.add_argument( '-i', '--inverse', action='store_true', help='inverse match country code', ) parser.add_argument( '--show', action='store_true', help='show brief info about input' ) parser.add_argument('-o', '--output', type=Path, help='output file to write') parser.add_argument( '-p', '--prefix', type=str, default='', help='write line prefix' ) parser.add_argument( '-s', '--suffix', type=str, default='\n', help='write line suffix' ) parser.set_defaults(func=main) subparser = parser.add_subparsers(dest='data_type', required=True) parser_geoip = subparser.add_parser( 'geoip', help='use geoip as input', formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser_geosite = subparser.add_parser( 'geosite', help='use geosite as input', formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser_geosite.add_argument( '--formatter', type=Path, help='format file for domain types', ) parser_geosite.add_argument( '--domain-types', nargs='+', type=int, help='select domain types (separated by space)', ) args = parser.parse_args() args.func(args)