1 #!/usr/bin/env python
 2 # coding=utf-8
 3 
 4 import re
 5 import sys
 6 import signal
 7 import codecs
 8 import publicsuffix
 9 import threading
10 import threadpool
11 from lib.data import kb
12 from lib.scan import ip_to_domain, get_icp_data, brute_subdomain,
13                      dns_zone_transfer, domain_to_ip
14 
15 MAX_DEPTH = 2
16 #PSL_FILE = publicsuffix.fetch()
17 PSL_FILE = codecs.open('public_suffix_list.dat', encoding='utf8')
18 P_DOMAIN = publicsuffix.PublicSuffixList(PSL_FILE)
19 
20 
21 def check_ip(ip):
22     compile_ip = re.compile('^((25[0-5]|2[0-4]d|[01]?dd?).){3}(25[0-5]|2[0-4]d|[01]?dd?)$')
23     return True if compile_ip.match(ip) else False
24 
25 
26 def monitor():
27     task_pool = threadpool.ThreadPool(10)
28     while len(kb.task_list_ip) or len(kb.task_list_domain) or threading.activeCount() > 1:
29         tip_msg = 'Ip task length: %s, Domain task length: %s.' % (len(kb.task_list_ip), len(kb.task_list_domain))
30         tip_msg += 'Ip result length: %s, Domain result length: %s.' % (len(kb.result_ip), len(kb.result_domain))
31         tip_msg += 'Workder Num: %s' % threading.activeCount()
32         print(tip_msg, end='
')
33 
34         if len(kb.task_list_ip):
35             ip, depth = kb.task_list_ip.pop()
36             if ip not in kb.task_done_ip and depth <= MAX_DEPTH:
37                 kb.task_done_ip.add(ip)
38                 wr = threadpool.WorkRequest(ip_to_domain, None, {'ip': ip, 'depth': depth})
39                 task_pool.putRequest(wr)
40 
41         if len(kb.task_list_domain):
42             domain, depth = kb.task_list_domain.pop()
43             root_domain = P_DOMAIN.get_public_suffix(domain)
44 
45             if root_domain not in kb.task_done_domain and depth <= MAX_DEPTH:
46                 kb.task_done_domain.add(root_domain)
47 
48                 wr = threadpool.WorkRequest(get_icp_data, None, {'root_domain': root_domain, 'depth': depth})
49                 task_pool.putRequest(wr)
50                 wr = threadpool.WorkRequest(brute_subdomain, None, {'root_domain': root_domain, 'depth': depth})
51                 task_pool.putRequest(wr)
52                 wr = threadpool.WorkRequest(dns_zone_transfer, None, {'root_domain': root_domain, 'depth': depth})
53                 task_pool.putRequest(wr)
54                 wr = threadpool.WorkRequest(domain_to_ip, None, {'domain': root_domain, 'depth': depth})
55                 task_pool.putRequest(wr)
56 
57             if domain != root_domain and domain not in kb.task_done_domain and depth <= MAX_DEPTH:
58                 kb.task_done_domain.add(domain)
59                 wr = threadpool.WorkRequest(domain_to_ip, None, {'domain': root_domain, 'depth': depth})
60                 task_pool.putRequest(wr)
61         try:
62             task_pool.poll()
63         except:
64             pass
65 
66     if task_pool.dismissedWorkers:
67         task_pool.joinAllDismissedWorkers()
68 
69 
70 def write_result():
71     print('Get %s domain and %s ip' % (len(kb.result_domain), len(kb.result_ip)))
72     with open('ip_result.txt', 'w') as f:
73         for ip in kb.result_ip:
74             f.write('%s
' % ip)
75 
76     with open('domain_result.txt', 'w') as f:
77         for domain in kb.result_domain:
78             f.write('%s
' % domain)
79 
80 
81 if __name__ == '__main__':
82     signal.signal(signal.SIGTERM, write_result)
83     signal.signal(signal.SIGINT, write_result)
84 
85     target = sys.argv[1]
86 
87     if check_ip(target):
88         kb.task_list_ip.add((target, 0))
89     else:
90         kb.task_list_domain.add((target, 0))
91 
92     print('Start scan with target %s ...' % target)
93 
94     t = threading.Thread(target=monitor, args=())
95     t.setDaemon(True)
96     t.start()
97     t.join()
98 
99     write_result()

$ git clone https://git@github.com:lewangbtcc/anti-XSS.git anti-XSS
$ cd anti-XSS
$ pip install -r requirements.txt