From 46134732fedfaf4933f0f8b620bae92aa1aa84ee Mon Sep 17 00:00:00 2001
From: dnomd343 <i@343.re>
Date: Tue, 2 Aug 2022 17:08:43 +0800
Subject: [PATCH] update: enhance check process

---
 Basis/Api.py        |  3 ++-
 Basis/Check.py      | 32 +++++++++++++++++++-------------
 Checker/Http.py     |  4 ++--
 Checker/__init__.py | 31 ++++++++++++++++++++++---------
 4 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/Basis/Api.py b/Basis/Api.py
index 1222425..2a25dfb 100644
--- a/Basis/Api.py
+++ b/Basis/Api.py
@@ -3,6 +3,7 @@
 
 import json
 from gevent import pywsgi
+from Checker import formatCheck
 from Basis.Logger import logging
 from Basis.Manager import Manager
 from Basis.Constant import Version
@@ -52,7 +53,7 @@ def createTask() -> Response:
         return genError('Invalid token')
 
     # TODO: format check and proxy list
-    checkList = request.json.get('check')
+    checkList = formatCheck(request.json.get('check'))
     proxyList = request.json.get('proxy')
 
     logging.debug('API create task -> check = %s | proxy = %s' % (checkList, proxyList))
diff --git a/Basis/Check.py b/Basis/Check.py
index 6fff046..b067b43 100644
--- a/Basis/Check.py
+++ b/Basis/Check.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+
 import copy
 import time
 from Checker import Checker
@@ -7,14 +8,9 @@ from Basis.Logger import logging
 from Builder import Builder, clientEntry
 
 
-def Check(taskId: str, taskInfo: dict) -> dict:
-    logging.info('[%s] Start checking process -> %s' % (taskId, taskInfo))
-    if taskInfo['type'] not in clientEntry:
-        logging.error('[%s] Unknown proxy type %s' % (taskId, taskInfo['type']))
-        raise RuntimeError('Unknown proxy type')
-
+def buildClient(taskId: str, taskInfo: dict) -> Builder:
     try:
-        client = Builder(
+        return Builder(
             proxyType = taskInfo['type'],
             proxyInfo = taskInfo['info'],
             bindAddr = '127.0.0.1',  # socks5 exposed host
@@ -23,25 +19,35 @@ def Check(taskId: str, taskInfo: dict) -> dict:
     except Exception as reason:
         logging.error('[%s] Client build error -> %s' % (taskId, reason))
         raise RuntimeError('Client build error')
-    logging.info('[%s] Client loaded successfully' % taskId)
 
+
+def waitClient(taskId: str, client: Builder):
     # TODO: wait port occupied (client.socksPort)
-    time.sleep(1)
+    time.sleep(1)  # TODO: simple delay for now
     if not client.status():  # client unexpected exit
         logging.warning('[%s] Client unexpected exit' % taskId)
         client.destroy()  # remove file and kill sub process
         logging.debug('[%s] Client output\n%s', (taskId, client.output))
         raise RuntimeError('Client unexpected exit')
 
-    checkResult = Checker(taskId, taskInfo['check'], {
+
+def Check(taskId: str, taskInfo: dict) -> dict:
+    logging.info('[%s] Start checking process -> %s' % (taskId, taskInfo))
+    if taskInfo['type'] not in clientEntry:  # unknown proxy type
+        logging.error('[%s] Unknown proxy type %s' % (taskId, taskInfo['type']))
+        raise RuntimeError('Unknown proxy type')
+    client = buildClient(taskId, taskInfo)  # build proxy client
+    logging.info('[%s] Client loaded successfully' % taskId)
+    waitClient(taskId, client)  # wait for the client to start
+    checkResult = Checker(taskId, taskInfo['check'], {  # start check process
         'addr': client.socksAddr,
         'port': client.socksPort,
     })
     logging.info('[%s] Client check result -> %s' % (taskId, checkResult))
-    client.destroy()  # clean up the client
+    client.destroy()  # clean up client
     taskInfo = copy.deepcopy(taskInfo)
-    taskInfo.pop('check')
+    taskInfo.pop('check')  # remove check items
     return {
         **taskInfo,
-        'result': checkResult,
+        'result': checkResult,  # add check result
     }
diff --git a/Checker/Http.py b/Checker/Http.py
index c62571c..d25f51b 100644
--- a/Checker/Http.py
+++ b/Checker/Http.py
@@ -28,7 +28,7 @@ def httpPing(taskId: str, url: str, socksAddr: str, socksPort: int, timeout: int
 
 
 def check(taskId: str, socksInfo: dict, options: dict) -> dict:
-    # TODO: multi check
-    return {  # TODO: just demo
+    # TODO: multi check (options -> times)
+    return {
         'delay': httpPing(taskId, options['url'], socksInfo['addr'], socksInfo['port'], options['timeout'])
     }
diff --git a/Checker/__init__.py b/Checker/__init__.py
index df894c3..d2501d4 100644
--- a/Checker/__init__.py
+++ b/Checker/__init__.py
@@ -1,18 +1,31 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+import logging
 
 from Checker import Http
 
-def Checker(taskId: str, checkInfo: dict, socksInfo: dict) -> dict:
-
-    # TODO: ignore checkInfo for now
+checkEntry = {
+    'http': Http.check
+}
 
-    httpRet = Http.check(taskId, socksInfo, {
-        'url': 'http://baidu.com/',
-        'timeout': 20,
-    })
+def formatCheck(rawInfo: list) -> dict:
+    # TODO: format check info
+    # TODO: rawInfo -> ['...', {'type': '...', ...}, ...]
     return {
-        'http': httpRet  # TODO: just check http delay for now
+        'http': {
+            'times': 3,
+            'url': 'http://baidu.com',
+            'timeout': 20,
+        }
     }
 
-    # TODO: return check result
+
+def Checker(taskId: str, checkInfo: dict, socksInfo: dict) -> dict:
+    diffItems = {x for x in checkInfo} - {x for x in checkEntry}
+    if len(diffItems) != 0:  # include unknown check items
+        logging.error('[%s] Unknown check items -> %s' % (taskId, diffItems))
+        raise RuntimeError('Unknown check items')
+    result = {}
+    for checkItem, checkOptions in checkInfo.items():
+        result[checkItem] = checkEntry[checkItem](taskId, socksInfo, checkOptions)
+    return result