加入终端颜色

Author: boy-hack

config.py

@@ -45,7 +45,4 @@
 
 # DEBUG
 DEBUG = True
-# EXCLUDE_PLUGINS = ['subdomain_found.py', 'jsonp.py', 'js_sensitive_content.py'
-#                    ]
-# INCLUDE_PLUGINS = ['sql_inject_bool.py']
-# THREAD_NUM = 1
+# EXCLUDE_PLUGINS = ['subdomain_found.py', 'jsonp.py', 'js_sensitive_content.py']

lib/data.py

@@ -3,15 +3,12 @@
 # @Time    : 2019/6/28 12:47 PM
 # @Author  : w8ay
 # @File    : data.py
-import logging
 
 from lib.datatype import _ThreaData
+from lib.log import LOGGER
 
-logging.basicConfig(level=logging.INFO,
-                    format='[%(asctime)s] %(levelname)s %(message)s',
-                    datefmt='%Y-%m-%d %H:%M:%S')
 
-logger = logging
+logger = LOGGER
 PATH = dict()  # 全局路径
 KB = dict()
 Share = _ThreaData()

lib/datatype.py

@@ -3,16 +3,11 @@
 # @Time    : 2019/6/29 12:27 AM
 # @Author  : w8ay
 # @File    : datatype.py
-import threading
 import logging
-from urllib.parse import urlparse
+import threading
 
 from lib.common import dataToStdout
 
-logging.basicConfig(level=logging.INFO,
-                    format='[%(asctime)s] %(levelname)s %(message)s',
-                    datefmt='%Y-%m-%d %H:%M:%S')
-
 
 class _ThreaData(object):
 

lib/log.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# @Time    : 2019/7/25 10:08 PM
+# @Author  : w8ay
+# @File    : log.py
+import logging
+
+FORMATTER = logging.Formatter("\r[%(asctime)s] [%(levelname)s] %(message)s", "%H:%M:%S")
+LOGGER = logging.getLogger('w13scan')
+LOGGER.setLevel(logging.INFO)
+LOGGER_HANDLER = logging.StreamHandler()
+LOGGER_HANDLER.setFormatter(FORMATTER)
+LOGGER.addHandler(LOGGER_HANDLER)

lib/option.py

@@ -9,6 +9,8 @@
 import time
 from queue import Queue
 
+from colorama import Fore
+
 from config import EXCLUDE_PLUGINS, INCLUDE_PLUGINS, SERVER_ADDR, DEBUG, INCLUDES, EXCLUDES, THREAD_NUM, LEVEL, TIMEOUT, \
     RETRY, PROXY_CONFIG, ACTIVE_SCAN, PROXY_CONFIG_BOOL
 from lib.common import dataToStdout
@@ -168,4 +170,5 @@ def banner():
     　 \　 /\　 /
     　　 ︶　 ︶
 '''
-    dataToStdout(_.format(version=VERSION, git=REPOSITORY))
+
+    dataToStdout(Fore.GREEN + _.format(version=VERSION, git=REPOSITORY))

lib/output.py

@@ -6,6 +6,8 @@
 
 from threading import Lock
 
+from colorama import Fore
+
 from lib.data import Share, KB
 
 
@@ -41,7 +43,7 @@ def success(self, url, plugin='unknown', **kw):
         report.update(kw)
         self.lock.acquire()
         self.collect.append(report)
-        self.log("[{}] ".format(report["plugin"]))
+        self.log("[{}] ".format(report["plugin"]), Fore.RED)
         del report["plugin"]
         raw = None
         if "raw" in report:
@@ -67,7 +69,7 @@ def success(self, url, plugin='unknown', **kw):
                 index += 1
         self.lock.release()
 
-    def log(self, msg):
+    def log(self, msg, color=Fore.YELLOW):
         width = KB["console_width"][0]
         outputs = []
         msgs = msg.split('\n')
@@ -80,7 +82,7 @@ def log(self, msg):
                 line = line[width:]
             outputs.append(line)
         for i in outputs:
-            Share.dataToStdout('\r' + i + ' ' * (width - len(i)) + '\n\r')
+            Share.dataToStdout('\r' + color + i + ' ' * (width - len(i)) + '\n\r')
 
     def output(self):
         '''

main.py

@@ -2,6 +2,8 @@
 import signal
 import threading
 
+from colorama import deinit, init as cinit, Style, Fore, Back
+
 from lib.baseproxy import AsyncMitmProxy
 from lib.cmdparse import cmd_line_parser
 from lib.controller import start
@@ -12,7 +14,9 @@
 def main():
     # init
     root = os.path.dirname(os.path.abspath(__file__))
+    cinit(autoreset=True)
     banner()
+
     cmdline = cmd_line_parser().__dict__
     init(root, cmdline)
     if conf["show_version"]:
@@ -30,6 +34,7 @@ def main():
         baseproxy.serve_forever()
     except KeyboardInterrupt:
         print("\n[*] User quit")
+        deinit()
         if not KB["is_win"]:
             os.kill(os.getpid(), signal.SIGHUP)
 

requirements.txt

@@ -2,11 +2,12 @@ asn1crypto==0.24.0
 certifi==2019.6.16
 cffi==1.12.3
 chardet==3.0.4
+colorama==0.4.1
 cryptography==2.7
 idna==2.8
 pycparser==2.19
 pyOpenSSL==18.0.0
 requests==2.22.0
 six==1.12.0
+tld==0.9.3
 urllib3==1.25.3
-tld==0.9.3