From 8a47ef2b9e3c00d56879b77e05324d45b06b8a77 Mon Sep 17 00:00:00 2001
From: antx-code <wkaifeng2007@163.com>
Date: Wed, 23 Mar 2022 17:53:08 +0800
Subject: [PATCH] update pocx requests exception error

---
 .coveragerc              |  7 +++++++
 .flake8                  |  5 +++++
 .gitignore               |  3 +++
 pocx/aio_poc.py          | 10 +++++-----
 pocx/basic_poc.py        | 10 +++++-----
 pocx/funcs/ceye.py       |  2 +-
 pocx/funcs/snow_flake.py |  3 +++
 7 files changed, 29 insertions(+), 11 deletions(-)
 create mode 100644 .coveragerc
 create mode 100644 .flake8

diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 0000000..a31d364
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,7 @@
+[run]
+
+source =
+    pocx
+    tests
+
+parallel = True
\ No newline at end of file
diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..af28862
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,5 @@
+[flake8]
+max-line-length = 88
+select = C,E,F,W,B,B9
+ignore = E203, E501, W503
+exclude = __init__.py
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 08d9e41..8a86391 100644
--- a/.gitignore
+++ b/.gitignore
@@ -131,3 +131,6 @@ dmypy.json
 dist
 .idea
 poetry.lock
+.coverage
+coverage.xml
+/tests/
diff --git a/pocx/aio_poc.py b/pocx/aio_poc.py
index 017f4ac..5b57e57 100644
--- a/pocx/aio_poc.py
+++ b/pocx/aio_poc.py
@@ -12,9 +12,11 @@ class AioPoc(metaclass=ABCMeta):
     def __init__(self) -> None:
         self.mode = 'Asynchronous Mode'
         self.name = "AioPoc"
-        self.example = ""
+        self.cve = ''
+        self.example = ''
         self.session = httpx.AsyncClient(verify=False)
         self.session.headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; rv:2.0.1) Gecko/20100101 Firefox/4.0.1"}
+        logger.info(f'Testing {self.name} with {self.mode}.')
 
     @logger.catch(level='ERROR')
     def set_headers(self, headers: dict = None):
@@ -25,8 +27,8 @@ def set_headers(self, headers: dict = None):
     async def aio_request(self, url: str, method: str = 'get', timeout: int = 10, **kwargs) -> httpx.Response:
         try:
             resp = await self.session.request(method, url, timeout=timeout, **kwargs)
-        except httpx.RequestError as e:
-            logger.error(f'[-] Run Poc[{self.name}] => {url} Error:\n\t{e}')
+        except Exception as e:
+            logger.error(f'[-] Run Poc [{self.cve} - {self.name}] Connection Error => {url} was not reachable.')
             resp = None
         return resp
 
@@ -57,8 +59,6 @@ async def poc(self, target: str) -> bool:
 
     @logger.catch(level='ERROR')
     def run(self, target):
-        logger.info(f'testing {self.name} with {self.mode}.')
-
         targets = [target] if isinstance(target, str) else target
 
         loop = asyncio.get_event_loop()
diff --git a/pocx/basic_poc.py b/pocx/basic_poc.py
index ad4fcce..533dda1 100644
--- a/pocx/basic_poc.py
+++ b/pocx/basic_poc.py
@@ -11,9 +11,11 @@ class BasicPoc(metaclass=ABCMeta):
     def __init__(self) -> None:
         self.mode = 'Synchronous Mode'
         self.name = "BasicPoc"
-        self.example = ""
+        self.cve = ''
+        self.example = ''
         self.session = httpx.Client(verify=False)
         self.session.headers = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; rv:2.0.1) Gecko/20100101 Firefox/4.0.1"}
+        logger.info(f'Testing {self.name} with {self.mode}.')
 
     @logger.catch(level='ERROR')
     def set_headers(self, headers: dict = None):
@@ -24,8 +26,8 @@ def set_headers(self, headers: dict = None):
     def request(self, url: str, method: str = 'get', timeout: int = 10, **kwargs) -> httpx.Response:
         try:
             resp = self.session.request(method, url, timeout=timeout, **kwargs)
-        except httpx.RequestError as e:
-            logger.error(f'[-] Run Poc[{self.name}] => {url} Error:\n\t{e}')
+        except Exception as e:
+            logger.error(f'[-] Run Poc [{self.cve} - {self.name}] Connection Error => {url} was not reachable.')
             resp = None
         return resp
 
@@ -56,8 +58,6 @@ def poc(self, target: str) -> bool:
 
     @logger.catch(level='ERROR')
     def run(self, target):
-        logger.info(f'testing {self.name} with {self.mode}.')
-
         targets = [target] if isinstance(target, str) else target
 
         for target in targets:
diff --git a/pocx/funcs/ceye.py b/pocx/funcs/ceye.py
index ae3c9ef..3db3807 100644
--- a/pocx/funcs/ceye.py
+++ b/pocx/funcs/ceye.py
@@ -39,7 +39,7 @@ def generate_payload_id(self):
         return str(self.id_worker.get_id())
 
     @logger.catch(level='ERROR')
-    def verify(self, pfilter: str, verify_type: str='dns'):
+    def verify(self, pfilter: str, verify_type: str = 'dns'):
         """
 
         Verify the payload.
diff --git a/pocx/funcs/snow_flake.py b/pocx/funcs/snow_flake.py
index cfe13a9..3aab98c 100644
--- a/pocx/funcs/snow_flake.py
+++ b/pocx/funcs/snow_flake.py
@@ -1,12 +1,14 @@
 # coding: utf-8
 import time
 
+
 class InvalidSystemClock(Exception):
     """
     时钟回拨异常
     """
     pass
 
+
 # 64位ID的划分
 WORKER_ID_BITS = 5
 DATACENTER_ID_BITS = 5
@@ -27,6 +29,7 @@ class InvalidSystemClock(Exception):
 # 开始时间截 (1995-03-19)
 TWEPOCH = 795542400000
 
+
 class IdWorker(object):
     """
     用于生成IDs