From 7c1ca5ffe07a35d85258eb5a188a5fe123019051 Mon Sep 17 00:00:00 2001
From: darkclip <darkclip@users.noreply.github.com>
Date: Wed, 19 Jul 2023 22:27:52 +0800
Subject: [PATCH] update

---
 raydata.py | 59 +++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 47 insertions(+), 12 deletions(-)

diff --git a/raydata.py b/raydata.py
index 9550f6db..c470ddf4 100755
--- a/raydata.py
+++ b/raydata.py
@@ -11,9 +11,10 @@ import ray_pb2
 
 
 def load_networks(
-    dat_files: typing.List[str], country_code: str
+    dat_files: typing.List[str], 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 = Path(file).resolve()
         if not file_path.exists() or not file_path.is_file():
@@ -22,8 +23,12 @@ def load_networks(
         with open(file_path, 'rb') as fp:
             data.ParseFromString(fp.read())
         for entry in data.entry:
-            if entry.country_code.lower() != country_code.lower():
-                continue
+            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)))
@@ -37,9 +42,13 @@ def load_networks(
 
 
 def load_domains(
-    dat_files: typing.List[str], country_code: str, domain_type: int | None = None
+    dat_files: typing.List[str],
+    countries: typing.List[str],
+    inverse: bool = False,
+    domain_type: int | None = None,
 ) -> typing.List[str]:
     domains = []
+    country_codes = [code.lower() for code in countries]
     for file in dat_files:
         file_path = Path(file).resolve()
         if not file_path.exists() or not file_path.is_file():
@@ -48,26 +57,40 @@ def load_domains(
         with open(file_path, 'rb') as fp:
             data.ParseFromString(fp.read())
         for entry in data.entry:
-            if entry.country_code.lower() != country_code.lower():
-                continue
+            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
                 if domain_type is None:
-                    domains.append(domain.value)
+                    value = domain.value
                 elif domain.type == domain_type:
+                    value = domain.value
+                if value:
                     domains.append(domain.value)
     return domains
 
 
 def main(args):
-    input = getattr(args, 'input', None)
+    input_files = getattr(args, 'input', None)
     output = getattr(args, 'output', None)
 
     data_list = []
     if getattr(args, 'data_type', None) == 'geoip':
-        data_list = load_networks(input, getattr(args, 'country', None))
+        data_list = load_networks(
+            input_files,
+            getattr(args, 'countries', None),
+            getattr(args, 'inverse', None),
+        )
     if getattr(args, 'data_type', None) == 'geosite':
         data_list = load_domains(
-            input, getattr(args, 'country', None), getattr(args, 'domain_type', None)
+            input_files,
+            getattr(args, 'countries', None),
+            getattr(args, 'inverse', None),
+            getattr(args, 'domain_type', None),
         )
     if len(data_list) == 0:
         print('No legal data.')
@@ -102,7 +125,20 @@ if __name__ == '__main__':
         default=[],
         help='input ray data files (separated by space)',
     )
-    parser.add_argument('-c', '--country', type=str, default='', help='country code')
+    parser.add_argument(
+        '-c',
+        '--countries',
+        nargs='+',
+        type=str,
+        default=[],
+        help='country codes (separated by space)',
+    )
+    parser.add_argument(
+        '-i',
+        '--inverse',
+        action='store_true',
+        help='inverse match country code',
+    )
     parser.add_argument('-o', '--output', type=Path, help='output file to write')
     parser.add_argument(
         '-p', '--prefix', type=str, default='', help='write line prefix'
@@ -126,7 +162,6 @@ if __name__ == '__main__':
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     parser_geosite.add_argument(
-        '-t',
         '--domain-type',
         type=int,
         help='geosite domain type',