import re
import requests
import time
import sys
import json
import pymongo
import threading

db_conn = pymongo.MongoClient('127.0.0.1',27017)
ls_db = db_conn.wangzuxian
ls_Waf = ls_db.waf

add_time = time.strftime('%Y-%m-%d',time.localtime(time.time()))

f  = open('cnt.txt','a+')

class waf_url(object):
   """"""
   #----------------------------------------------------------------------
   def __init__(self,url):
       """Constructor"""
       super(waf_url, self).__init__()
       self.url = url
       self.waf_name = ""
       self.headers = []
       
   #----------------------------------------------------------------------
   def waf_rules(self):
       """"""
       with open('waf_rules.txt') as waf_data:
           for data in waf_data:
               if re.match("\[.*?\]|^;", data) or not data.split():
                   continue
               name, location, key, reg = data.strip().split("|", 3)
               try:
                   if key in self.headers and (re.search(reg, self.headers[key], re.I)): 
                       #print key
                       rule_type, rule_name = name.split(':')
                       ls_Waf.update({"URL":self.url},
                                     {"$set": {rule_type : rule_name,'add_time':add_time}},
                                     upsert = True)  
                   else:
                       continue             
               except Exception,e:
                   print str(e)
                   pass
   #----------------------------------------------------------------------
   def waf_etag(self):
       """"""
       for k in self.headers.keys():
           if 'ETag'.upper() == k.upper():
               print self.url + '------Etag'  
               ls_Waf.update({"URL":self.url},
                                     {"$set": {'ETag':'True','add_time':add_time}},
                                                upsert = True)  
           if re.search('cdn'.upper(), self.headers[k].upper()):
               ls_Waf.update({"URL":self.url},
                                         {"$set": {'CDN':self.headers[k],'add_time':add_time}},
                                                 upsert = True)                
               
           else:
               continue
       

   #----------------------------------------------------------------------
   def waf_rule_headers(self):
       """"""
       try:
           r = requests.get(self.url)
           requests.adapters.DEFAULT_RETRIES = 5 
           #print self.url
           self.headers = r.headers
           self.waf_rules()
           self.waf_etag()
       except Exception,e:
           print str(e)
           f.write(self.url)
           pass
         
def start_waf():
   """"""
   with open('kehu_ok.txt') as shang:
       for i in shang:        
           i = i.strip('\r').strip('\n')     
           identify_waf = waf_url(i)
           identify_waf.waf_rule_headers()
   

for num in range(100):
   t = threading.Thread(target= start_waf())
   t.setDaemon(True)
   t.start()

   
f.close()