add -k args

Author: kingkaki

README.md

@@ -12,6 +12,11 @@ Example: python ctf-wscan.py http://ctf.test.com
 
 ![](1.gif)
 
+# 新增-k参数
+```
+Useage : python ctf-wscan.py [website url] [key words ……]
+```
+由于觉得添加关键字这个功能还是蛮常用的，于是添加了个`-k`参数，可以在命令行中添加参数
 
 
 # 一些设置

config.py

@@ -2,7 +2,7 @@
 # @Author: King kaki
 # @Date:   2018-07-30 12:40:36
 # @Last Modified by:   kingkk
-# @Last Modified time: 2018-08-12 09:54:54
+# @Last Modified time: 2018-08-18 15:59:05
 
 # 关键字
 KEY_WORDS = ['flag','ctf', 'admin']

ctf-wscan.py

@@ -2,16 +2,22 @@
 # @Author: King kaki
 # @Date:   2018-07-30 12:37:36
 # @Last Modified by:   kingkk
-# @Last Modified time: 2018-08-12 09:58:53
+# @Last Modified time: 2018-08-18 15:33:01
 
 # url = 'http://localhost:80/'
 # url = 'http://ctf5.shiyanbar.com/web/'
 import sys
-
+import argparse
 from lib.init import Init
 
+
 def main():
-	scan = Init(sys.argv)
+	parser = argparse.ArgumentParser()
+	parser.add_argument('url', help="The website to be scanned", type=str)
+	parser.add_argument('-k', '--keys', dest="key_words", nargs='+', help="Keys words to extend scan", type=str, default="")
+	args = parser.parse_args()
+	
+	scan = Init(args)
 	scan.start()
 
 if __name__ == '__main__':

dict/default.txt

@@ -42,12 +42,12 @@ _viminfo
 %3f.back
 phpinfo.php
 robots.txt
-test.php
 .htaccess
 .bash_history
 .svn/
 .git/
 .index.php.swp
+index.php.swp
 index.php.bak
 .index.php~
 index.php.bak_Edietplus
@@ -148,7 +148,6 @@ login.php
 register.php
 upload.php
 home.php
-test.php
 log.php
 logs.php
 config.php

dict/ext.txt

@@ -1,3 +1,4 @@
+$
 $.txt
 $.html
 $/
@@ -11,6 +12,7 @@ $.gz
 $.tar
 $~
 $.swp
+.$.swp
 $.swp~
 $.save
 $.back

lib/generatedict.py

@@ -1,17 +1,24 @@
 # -*- coding: utf-8 -*-
 # @Author: King kaki
 # @Date:   2018-07-30 14:24:51
-# @Last Modified by:   King kaki
-# @Last Modified time: 2018-07-30 15:08:32
+# @Last Modified by:   kingkk
+# @Last Modified time: 2018-08-18 16:00:09
 
 from config import *
 
 import re
 
 class GenerateDcit:
-	def __init__(self):
+	def __init__(self, keywords):
 		self.exts = self._getexts()
-		self.keywords = KEY_WORDS
+
+		self.keywords = []
+
+		if KEY_WORDS:
+			self.keywords += KEY_WORDS
+
+		if keywords:
+			self.keywords += keywords
 
 
 	def _getexts(self):

lib/init.py

@@ -2,19 +2,18 @@
 # @Author: kingkk
 # @Date:   2018-08-11 19:32:38
 # @Last Modified by:   kingkk
-# @Last Modified time: 2018-08-12 10:27:39
+# @Last Modified time: 2018-08-18 15:58:15
 import sys
 from config import * 
 from lib.log import Log
 from lib.scan import Scan
+from lib.generatedict import GenerateDcit
+
 
 class Init:
-	def __init__(self, argv):
-		if len(argv) != 2:
-			print('Params Error!')
-			self.help()
-		self.url = str(argv[1])
-		self.url = self.init_url(self.url)
+	def __init__(self, args):	
+		self.url = self.init_url(str(args.url))
+		self.keywords = args.key_words
 
 
 	def help(self):
@@ -37,11 +36,28 @@ def init_url(self, url):
 
 		return url
 
+	def get_files(self):
+		# 获取默认扫描列表
+		with open('dict/default.txt') as f:
+			files = f.readlines()
+		
+
+		#生成关键字字典
+		if KEY_WORDS or args.KEY_WORDS:
+			g = GenerateDcit(self.keywords)
+			for i in g.generate():
+				files.append(i)
+
+
+		files = (file.strip() for file in files)
+		return files
+
 	def start(self):
 		threadlist = []
 		loglist = {}
+		files = self.get_files()
 		for i in range(NUMBER_OF_THREAD):
-			threadlist.append(Scan(self.url, loglist))
+			threadlist.append(Scan(self.url, loglist, files))
 		for t in threadlist:
 			t.start()
 		for t in threadlist:

lib/scan.py

@@ -2,58 +2,41 @@
 # @Author: King kaki
 # @Date:   2018-07-30 13:18:58
 # @Last Modified by:   kingkk
-# @Last Modified time: 2018-08-12 09:55:23
+# @Last Modified time: 2018-08-18 15:59:23
 
 import sys
 import re
 import threading
-
 from config import *
-from lib.generatedict import GenerateDcit
 
 import requests
 
 
-def setting():
-	# 获取请求方式
-	if REQUEST_METHOD == 1:
-		req = requests.head
-	elif REQUEST_METHOD == 2:
-		req = requests.get
-
-	# 获取默认扫描列表
-	with open('dict/default.txt') as f:
-		files = f.readlines()
-	
-
-	#生成关键字字典
-	if KEY_WORDS:
-		g = GenerateDcit()
-		for i in g.generate():
-			files.append(i)
-
-
-	files = (file.strip() for file in files)
-
-	return req, files
-
-req, files = setting()
-
 class Scan(threading.Thread):
-	def __init__(self, url, log):
+	def __init__(self, url, log, files):
 		threading.Thread.__init__(self)
 		self.url = url
 		self.log = log
+		self.req = self.setting()
+		self.files = files
 
 	def run(self):
-		for file in files:
+		for file in self.files:
 			try:
-				r = req(self.url+file, timeout=TIME_OUT)
+				r = self.req(self.url+file, timeout=TIME_OUT)
 			except:
 				continue
 			with threading.Lock():
 				self.display(r, file)
 
+	def setting(self):
+		# 获取请求方式
+		if REQUEST_METHOD == 1:
+			req = requests.head
+		elif REQUEST_METHOD == 2:
+			req = requests.get
+		return req
+
 
 	def display(self, r, file):
 		if r.status_code not in INVALID_CODE: