feat: filter proxies by country via ?area=CN (#84, #177) (#233)

* feat: support fetching multiple proxies via /random?count=N

- redis: add RedisClient.randoms(count, ...) which prefers max-score proxies and falls back to rank, returning a de-duplicated random sample

- server: /random accepts a count parameter and returns N random proxies (one per line); count<=1 keeps the original single-proxy behavior; works with key too

- README: document the new count parameter

Closes #194

* feat: filter proxies by country via ?area=CN on /random and /all

- utils/geo: add get_country_iso() backed by the bundled GeoLite2 db (previously an unused dependency); fails safe to None if unavailable

- server: /random and /all accept an area parameter (country iso code) to return only proxies from that country; combines with count and key

- README: document the area parameter

Closes #84, Closes #177

Author: Germey

README.md

@@ -222,6 +222,27 @@ get random proxy 116.196.115.209:8080
 
 可以看到成功获取了代理，并请求 httpbin.org 验证了代理的可用性。
 
+### 获取多个代理
+
+如果一次需要多个代理，可以给 `/random` 接口传入 `count` 参数，一次返回多个随机代理（每行一个）：
+
+```
+GET http://localhost:5555/random?count=5
+```
+
+`count` 不传或为 1 时行为不变，仍返回单个代理；`count` 大于可用数量时返回全部可用代理。也可与 `key` 参数组合使用。
+
+### 按地区（国家）筛选代理
+
+可以给 `/random` 和 `/all` 接口传入 `area` 参数，按代理 IP 所属国家筛选（ISO 国家码，大小写不敏感），例如只获取国内（中国）代理：
+
+```
+GET http://localhost:5555/random?area=CN
+GET http://localhost:5555/all?area=CN
+```
+
+国家信息由内置的 GeoLite2 离线库解析，无法解析归属地的代理会被排除。`area` 可与 `count`、`key` 参数组合使用。
+
 ## 可配置项
 
 代理池可以通过设置环境变量来配置一些参数。

proxypool/processors/server.py

@@ -5,6 +5,8 @@
 from proxypool.storages.redis import RedisClient
 from proxypool.setting import API_HOST, API_PORT, API_THREADED, API_KEY, IS_DEV, PROXY_RAND_KEY_DEGRADED
 import functools
+from random import choice, sample
+from proxypool.utils.geo import get_country_iso
 
 __all__ = ['app']
 
@@ -58,6 +60,20 @@ def get_request_key():
     return key
 
 
+def filter_proxies_by_area(proxies, area):
+    """
+    filter proxies by country iso code (e.g. 'CN', 'US'), case-insensitive;
+    proxies whose country cannot be resolved are excluded
+    :param proxies: list of Proxy
+    :param area: country iso code, or falsy to skip filtering
+    :return: filtered list of Proxy
+    """
+    if not area:
+        return proxies
+    area = area.upper()
+    return [proxy for proxy in proxies if get_country_iso(proxy.host) == area]
+
+
 @app.route('/')
 @auth_required
 def index():
@@ -74,11 +90,37 @@ def get_proxy():
     """
     get a random proxy, can query the specific sub-pool according the (redis) key
     if PROXY_RAND_KEY_DEGRADED is set to True, will get a universal random proxy if no proxy found in the sub-pool
+    can pass a `count` parameter to get multiple random proxies at once
+    can pass an `area` parameter to only get proxies from a country (iso code, e.g. CN)
     :return: get a random proxy
     """
     key = get_request_key()
+    count = request.args.get('count', type=int)
+    area = request.args.get('area')
     conn = get_conn()
     # return conn.random(key).string() if key else conn.random().string()
+    if area:
+        # area filtering needs the candidate set first, then filter by country
+        candidates = conn.all(key) if key else conn.all()
+        candidates = filter_proxies_by_area(candidates, area)
+        if not candidates and key and PROXY_RAND_KEY_DEGRADED:
+            candidates = filter_proxies_by_area(conn.all(), area)
+        if not candidates:
+            raise PoolEmptyException
+        if count and count > 1:
+            count = min(count, len(candidates))
+            return '\n'.join(proxy.string() for proxy in sample(candidates, count))
+        return choice(candidates).string()
+    if count and count > 1:
+        # return multiple random proxies, one per line
+        try:
+            proxies = conn.randoms(count, key) if key else conn.randoms(count)
+        except PoolEmptyException:
+            if key and PROXY_RAND_KEY_DEGRADED:
+                proxies = conn.randoms(count)
+            else:
+                raise
+        return '\n'.join(proxy.string() for proxy in proxies)
     if key:
         try:
             return conn.random(key).string()
@@ -92,13 +134,15 @@ def get_proxy():
 @auth_required
 def get_proxy_all():
     """
-    get a random proxy
-    :return: get a random proxy
+    get all proxies, optionally filtered by `area` (country iso code, e.g. CN)
+    :return: all proxies
     """
     key = get_request_key()
+    area = request.args.get('area')
 
     conn = get_conn()
     proxies = conn.all(key) if key else conn.all()
+    proxies = filter_proxies_by_area(proxies, area)
     proxies_string = ''
     if proxies:
         for proxy in proxies:

proxypool/storages/redis.py

@@ -3,7 +3,7 @@
 from proxypool.schemas.proxy import Proxy
 from proxypool.setting import REDIS_CONNECTION_STRING, REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, REDIS_DB, REDIS_KEY, PROXY_SCORE_MAX, PROXY_SCORE_MIN, \
     PROXY_SCORE_INIT
-from random import choice
+from random import choice, sample
 from typing import List
 from loguru import logger
 from proxypool.utils.proxy import is_valid_proxy, convert_proxy_or_proxies
@@ -70,6 +70,27 @@ def random(self, redis_key=REDIS_KEY, proxy_score_min=PROXY_SCORE_MIN, proxy_sco
         # else raise error
         raise PoolEmptyException
 
+    def randoms(self, count, redis_key=REDIS_KEY, proxy_score_min=PROXY_SCORE_MIN, proxy_score_max=PROXY_SCORE_MAX) -> List[Proxy]:
+        """
+        get a batch of random proxies
+        firstly try to get proxies with max score,
+        if not enough, get proxies by rank (score from high to low)
+        if none exists, raise error
+        :param count: number of proxies to return
+        :return: list of proxies
+        """
+        # try to get proxies with max score first
+        proxies = self.db.zrangebyscore(
+            redis_key, proxy_score_max, proxy_score_max)
+        if len(proxies) < count:
+            # not enough max-score proxies, fall back to all proxies by rank
+            proxies = self.db.zrevrangebyscore(
+                redis_key, proxy_score_max, proxy_score_min)
+        if not proxies:
+            raise PoolEmptyException
+        count = min(count, len(proxies))
+        return convert_proxy_or_proxies(sample(proxies, count))
+
     def decrease(self, proxy: Proxy, redis_key=REDIS_KEY, proxy_score_min=PROXY_SCORE_MIN) -> int:
         """
         decrease score of proxy, if small than PROXY_SCORE_MIN, delete it

proxypool/utils/geo.py

@@ -0,0 +1,35 @@
+from loguru import logger
+
+# geolite2 provides an offline IP -> country database (bundled with the
+# maxminddb_geolite2 dependency). loading it can fail if the optional
+# dependency is missing, so degrade gracefully and disable area filtering.
+try:
+    from geolite2 import geolite2
+
+    _reader = geolite2.reader()
+except Exception as e:  # pragma: no cover
+    _reader = None
+    logger.warning(f'geolite2 is unavailable, area filtering disabled: {e}')
+
+
+def get_country_iso(ip):
+    """
+    look up the ISO country code (e.g. 'CN', 'US') for an ip address
+    :param ip: ip address string
+    :return: uppercase iso code, or None if unknown/unavailable
+    """
+    if _reader is None:
+        return None
+    try:
+        record = _reader.get(ip)
+    except Exception:
+        return None
+    if not record:
+        return None
+    country = record.get('country') or record.get('registered_country') or {}
+    return country.get('iso_code')
+
+
+if __name__ == '__main__':
+    print('8.8.8.8', get_country_iso('8.8.8.8'))
+    print('114.114.114.114', get_country_iso('114.114.114.114'))