Merge pull request #302 from Guovin/dev

Dev

.github/workflows/main.yml

@@ -115,6 +115,12 @@ jobs:
           elif [[ -f "output/result.log" ]]; then
             git add -f "output/result.log"
           fi
+          if [[ -f "updates/fofa/fofa_hotel_region_result.pkl" ]]; then
+            git add -f "updates/fofa/fofa_hotel_region_result.pkl"
+          fi
+          if [[ -f "updates/fofa/fofa_multicast_region_result.pkl" ]]; then
+            git add -f "updates/fofa/fofa_multicast_region_result.pkl"
+          fi
           if ! git diff --staged --quiet; then
             git commit -m "Github Action Auto Updated"
             git push --force

Pipfile

@@ -25,7 +25,6 @@ opencc-python-reimplemented = "*"
 fake-useragent = "*"
 gunicorn = "*"
 pillow = "*"
-rapidfuzz = "*"
 
 [requires]
 python_version = "3.8"

main.py

@@ -93,7 +93,11 @@ async def visit_page(self, channel_names=None):
                 continue
             if config.getboolean("Settings", setting):
                 if setting == "open_subscribe":
-                    subscribe_urls = config.get("Settings", "subscribe_urls").split(",")
+                    subscribe_urls = [
+                        url.strip()
+                        for url in config.get("Settings", "subscribe_urls").split(",")
+                        if url.strip()
+                    ]
                     task = asyncio.create_task(
                         task_func(subscribe_urls, callback=self.update_progress)
                     )

tkinter_ui/hotel.py

@@ -71,7 +71,7 @@ def init_ui(self, root):
         self.region_list_label.pack(side=tk.LEFT, padx=4, pady=8)
         regions = ["全部"] + list(getattr(fofa_map, "region_url").keys())
         region_selected_values = [
-            value
+            value.strip()
             for value in config.get("Settings", "hotel_region_list").split(",")
             if value.strip()
         ]

tkinter_ui/multicast.py

@@ -86,7 +86,7 @@ def init_ui(self, root):
             regions.remove("全部")
         regions.insert(0, "全部")
         region_selected_values = [
-            value
+            value.strip()
             for value in config.get("Settings", "multicast_region_list").split(",")
             if value.strip()
         ]

tkinter_ui/select_combobox.py

@@ -24,7 +24,7 @@ def on_select(self, event):
 
     def on_text_change(self, event):
         text_value = self.get().strip()
-        value_list = [value.strip() for value in text_value.split(",")]
+        value_list = [value.strip() for value in text_value.split(",") if value.strip()]
         self.selected_values = [
             value for value in self.selected_values if value in value_list
         ]

tkinter_ui/tkinter_ui.spec

@@ -10,6 +10,8 @@ a = Analysis(
         ('../config/rtp', 'config/rtp'),
         ('../updates/multicast/multicast_map.json', 'updates/multicast'),
         ('../updates/multicast/multicast_region_result.json', 'updates/multicast'),
+        ('../updates/fofa/fofa_hotel_region_result.pkl', 'updates/fofa'),
+        ('../updates/fofa/fofa_multicast_region_result.pkl', 'updates/fofa'),
         ('../static/images/favicon.ico', 'static/images'),
         ('../static/images/appreciate.jpg', 'static/images'),
         ('../static/images/settings_icon.png', 'static/images'),

updates/fofa/request.py

@@ -1,8 +1,8 @@
-from utils.config import config
+from utils.config import config, resource_path
 from tqdm.asyncio import tqdm_asyncio
 from time import time
 from requests import get
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import updates.fofa.fofa_map as fofa_map
 from driver.setup import setup_driver
 import re
@@ -12,6 +12,8 @@
 from updates.proxy import get_proxy, get_proxy_next
 from requests_custom.utils import get_source_requests, close_session
 from collections import defaultdict
+import pickle
+import threading
 
 timeout = 10
 
@@ -20,7 +22,11 @@ def get_fofa_urls_from_region_list():
     """
     Get the FOFA url from region
     """
-    region_list = config.get("Settings", "hotel_region_list").split(",")
+    region_list = [
+        region.strip()
+        for region in config.get("Settings", "hotel_region_list").split(",")
+        if region.strip()
+    ]
     urls = []
     region_url = getattr(fofa_map, "region_url")
     if "all" in region_list or "ALL" in region_list or "全部" in region_list:
@@ -32,6 +38,30 @@ def get_fofa_urls_from_region_list():
     return urls
 
 
+def update_fofa_region_result_tmp(result, multicast=False):
+    """
+    Update fofa region result tmp
+    """
+    with open(
+        resource_path(
+            f"updates/fofa/fofa_{'multicast' if multicast else 'hotel'}_region_result.pkl"
+        ),
+        "wb",
+    ) as file:
+        pickle.dump(result, file)
+
+
+def get_fofa_region_result_tmp(multicast: False):
+    with open(
+        resource_path(
+            f"updates/fofa/fofa_{'multicast' if multicast else 'hotel'}_region_result.pkl"
+        ),
+        "rb",
+    ) as file:
+        result = pickle.load(file)
+        return result
+
+
 async def get_channels_by_fofa(urls=None, multicast=False, callback=None):
     """
     Get the channel by FOFA
@@ -57,9 +87,12 @@ async def get_channels_by_fofa(urls=None, multicast=False, callback=None):
     if open_proxy:
         test_url = fofa_urls[0][0] if multicast else fofa_urls[0]
         proxy = await get_proxy(test_url, best=True, with_test=True)
+    cancel_event = threading.Event()
 
     def process_fofa_channels(fofa_info):
-        nonlocal proxy, fofa_urls_len, open_driver, open_sort
+        nonlocal proxy, fofa_urls_len, open_driver, open_sort, cancel_event
+        if cancel_event.is_set():
+            return {}
         fofa_url = fofa_info[0] if multicast else fofa_info
         results = defaultdict(lambda: defaultdict(list))
         try:
@@ -79,6 +112,9 @@ def process_fofa_channels(fofa_info):
                 page_source = retry_func(
                     lambda: get_source_requests(fofa_url), name=fofa_url
                 )
+            if "资源访问每天限制" in page_source:
+                cancel_event.set()
+                raise ValueError("Limited access to fofa page")
             fofa_source = re.sub(r"<!--.*?-->", "", page_source, flags=re.DOTALL)
             urls = set(re.findall(r"https?://[\w\.-]+:\d+", fofa_source))
             if multicast:
@@ -94,6 +130,9 @@ def process_fofa_channels(fofa_info):
                     ]
                     for future in futures:
                         results = merge_objects(results, future.result())
+            return results
+        except ValueError as e:
+            raise e
         except Exception as e:
             print(e)
         finally:
@@ -107,15 +146,32 @@ def process_fofa_channels(fofa_info):
                     f"正在获取Fofa{mode_name}源, 剩余{remain}个查询地址待获取, 预计剩余时间: {get_pbar_remaining(n=pbar.n, total=pbar.total, start_time=start_time)}",
                     int((pbar.n / fofa_urls_len) * 100),
                 )
-            return results
 
     max_workers = 3 if open_driver else 10
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
         futures = [
             executor.submit(process_fofa_channels, fofa_url) for fofa_url in fofa_urls
         ]
-        for future in futures:
-            fofa_results = merge_objects(fofa_results, future.result())
+        try:
+            for future in as_completed(futures):
+                fofa_results = merge_objects(fofa_results, future.result())
+        except ValueError as e:
+            if "Limited access to fofa page" in str(e):
+                for future in futures:
+                    future.cancel()
+                fofa_results = {}
+    if fofa_results:
+        update_fofa_region_result_tmp(fofa_results, multicast=multicast)
+    else:
+        fofa_results = get_fofa_region_result_tmp(multicast=multicast)
+        print(fofa_results)
+        pbar.n = fofa_urls_len
+        pbar.update(0)
+        if callback:
+            callback(
+                f"正在获取Fofa{mode_name}源",
+                100,
+            )
     if not open_driver:
         close_session()
     pbar.close()

updates/hotel/request.py

@@ -35,7 +35,11 @@ async def get_channels_by_hotel(callback=None):
     open_proxy = config.getboolean("Settings", "open_proxy")
     open_driver = config.getboolean("Settings", "open_driver")
     page_num = config.getint("Settings", "hotel_page_num")
-    region_list = config.get("Settings", "hotel_region_list").split(",")
+    region_list = [
+        region.strip()
+        for region in config.get("Settings", "hotel_region_list").split(",")
+        if region.strip()
+    ]
     if "all" in region_list or "ALL" in region_list or "全部" in region_list:
         fofa_region_name_list = list(getattr(fofa_map, "region_url").keys())
         region_list = fofa_region_name_list

updates/multicast/update_tmp.py

@@ -103,7 +103,11 @@ def get_multicast_region_result_by_rtp_txt(callback=None):
     Get multicast region result by rtp txt
     """
     rtp_path = resource_path("config/rtp")
-    config_region_list = set(config.get("Settings", "multicast_region_list").split(","))
+    config_region_list = set(
+        region.strip()
+        for region in config.get("Settings", "multicast_region_list").split(",")
+        if region.strip()
+    )
     rtp_file_list = [
         filename.rsplit(".", 1)[0]
         for filename in os.listdir(rtp_path)

updates/subscribe/request.py

@@ -27,7 +27,7 @@ async def get_channels_by_subscribe_urls(
     subscribe_results = {}
     pattern = r"^(.*?),(?!#genre#)(.*?)$"
     subscribe_urls = [
-        url
+        url.strip()
         for url in config.get("Settings", "subscribe_urls").split(",")
         if url.strip()
     ]