From fddc472ad5c47024876fcd8a7f3a1bcd2607fc20 Mon Sep 17 00:00:00 2001
From: g1879 <g1879@qq.com>
Date: Sat, 29 Jun 2024 17:22:22 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96auto=5Fport()=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 DrissionPage/_configs/chromium_options.py  |  2 +-
 DrissionPage/_configs/chromium_options.pyi |  2 +-
 DrissionPage/_functions/tools.py           | 49 ++++++++++++----------
 DrissionPage/_functions/tools.pyi          |  4 +-
 4 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/DrissionPage/_configs/chromium_options.py b/DrissionPage/_configs/chromium_options.py
index 03f1152..994639d 100644
--- a/DrissionPage/_configs/chromium_options.py
+++ b/DrissionPage/_configs/chromium_options.py
@@ -512,7 +512,7 @@ class ChromiumOptions(object):
         :return: 当前对象
         """
         if on_off:
-            self._auto_port = scope if scope else True
+            self._auto_port = scope if scope else (9600, 19600)
             if tmp_path:
                 self._tmp_path = str(tmp_path)
         else:
diff --git a/DrissionPage/_configs/chromium_options.pyi b/DrissionPage/_configs/chromium_options.pyi
index 56fca1e..42013b5 100644
--- a/DrissionPage/_configs/chromium_options.pyi
+++ b/DrissionPage/_configs/chromium_options.pyi
@@ -27,7 +27,7 @@ class ChromiumOptions(object):
     _flags: dict = ...
     _prefs_to_del: list = ...
     clear_file_flags: bool = ...
-    _auto_port: bool = ...
+    _auto_port: Union[Tuple[int, int], False] = ...
     _system_user_path: bool = ...
     _existing_only: bool = ...
     _retry_times: int = ...
diff --git a/DrissionPage/_functions/tools.py b/DrissionPage/_functions/tools.py
index 2486718..4f6b13a 100644
--- a/DrissionPage/_functions/tools.py
+++ b/DrissionPage/_functions/tools.py
@@ -8,7 +8,7 @@
 from pathlib import Path
 from platform import system
 from shutil import rmtree
-from tempfile import gettempdir, TemporaryDirectory
+from tempfile import gettempdir
 from threading import Lock
 from time import perf_counter, sleep
 
@@ -18,8 +18,10 @@ from ..errors import (ContextLostError, ElementLostError, CDPError, PageDisconne
 
 
 class PortFinder(object):
-    used_port = {}
+    used_port = set()
+    prev_time = None
     lock = Lock()
+    checked_paths = set()
 
     def __init__(self, path=None):
         """
@@ -28,34 +30,39 @@ class PortFinder(object):
         tmp = Path(path) if path else Path(gettempdir()) / 'DrissionPage'
         self.tmp_dir = tmp / 'UserTempFolder'
         self.tmp_dir.mkdir(parents=True, exist_ok=True)
-        if not PortFinder.used_port:
-            clean_folder(self.tmp_dir)
+        if str(self.tmp_dir.absolute()) not in PortFinder.checked_paths:
+            for i in self.tmp_dir.iterdir():
+                if i.is_dir() and i.stem.startswith('AutoPort') and not port_is_using('127.0.0.1', i.name[8:]):
+                    rmtree(i, ignore_errors=True)
+            PortFinder.checked_paths.add(str(self.tmp_dir.absolute()))
 
     def get_port(self, scope=None):
         """查找一个可用端口
         :param scope: 指定端口范围，不含最后的数字，为None则使用[9600-19600)
         :return: 可以使用的端口和用户文件夹路径组成的元组
         """
+        from random import randint
         with PortFinder.lock:
+            if PortFinder.prev_time and perf_counter() - PortFinder.prev_time > 30:
+                PortFinder.used_port.clear()
+            PortFinder.prev_time = perf_counter()
             if scope in (True, None):
                 scope = (9600, 19600)
-            for i in range(scope[0], scope[1]):
-                if i in PortFinder.used_port:
+            msx_times = scope[1] - scope[0]
+            times = 0
+            while times < msx_times:
+                times += 1
+                port = randint(*scope)
+                if port in PortFinder.used_port or port_is_using('127.0.0.1', port):
                     continue
-                elif port_is_using('127.0.0.1', i):
-                    PortFinder.used_port[i] = None
-                    continue
-                path = TemporaryDirectory(dir=self.tmp_dir).name
-                PortFinder.used_port[i] = path
-                return i, path
-
-            for i in range(scope[0], scope[1]):
-                if port_is_using('127.0.0.1', i):
-                    continue
-                rmtree(PortFinder.used_port[i], ignore_errors=True)
-                return i, TemporaryDirectory(dir=self.tmp_dir).name
-
-        raise OSError('未找到可用端口。')
+                path = self.tmp_dir / f'AutoPort{port}'
+                if path.exists():
+                    try:
+                        rmtree(path)
+                    except:
+                        continue
+                return port, str(path)
+            raise OSError('未找到可用端口。')
 
 
 def port_is_using(ip, port):
@@ -95,7 +102,7 @@ def show_or_hide_browser(page, hide=True):
     :param hide: 是否隐藏
     :return: None
     """
-    if not page.address.startswith(('127.0.0.1', 'localhost')):
+    if not page.browser.address.startswith(('127.0.0.1', 'localhost')):
         return
 
     if system().lower() != 'windows':
diff --git a/DrissionPage/_functions/tools.pyi b/DrissionPage/_functions/tools.pyi
index a6fc535..4fa536b 100644
--- a/DrissionPage/_functions/tools.pyi
+++ b/DrissionPage/_functions/tools.pyi
@@ -14,9 +14,11 @@ from .._pages.chromium_base import ChromiumBase
 
 
 class PortFinder(object):
-    used_port: dict = ...
+    used_port: set = ...
+    prev_time: float = ...
     lock: Lock = ...
     tmp_dir: Path = ...
+    checked_paths: set = ...
 
     def __init__(self, path: Union[str, Path] = None): ...