返回
顶部

渗透测试过程中,经常需要编写一些python小工具

但是单线程太慢了,多线程可以极快地提高效率

我们可以直接通过一个例子来学习多线程脚本的编写

下面是一个验证域名是否有效的python脚本

#!/usr/bin/python

import Queue
import threading
import time
import os
import sys
import socket
exitFlag = False
f = open(sys.argv[1])
marylines = f.readlines()
f.close()
countttter = threading.Lock()
countttterssss=0
tottttt = len(marylines)
class myThread (threading.Thread):
   def __init__(self, threadID, name, q):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = name
        self.q = q
   def run(self):
        print "Starting " + self.name
        process_data(self.name, self.q)
        print "Exiting " + self.name

def process_data(threadName, q):
   while not exitFlag:
        queueLock.acquire()
        if not workQueue.empty():
            data = q.get()
            queueLock.release()
            print "%s processing %s" % (threadName, data)
            #print data
            try:
                socket.gethostbyname(data.strip())
                filelock.acquire()
                ssssssssssssf = open(sys.argv[3],'a')
                ssssssssssssf.write(data)
                ssssssssssssf.close()
                filelock.release()

            except socket.gaierror:
                print "unable to get address for: ", data
            countttter.acquire()
            global countttterssss
            countttterssss = countttterssss + 1
            print  str(countttterssss) + '-----------------------------' + str(tottttt) + '\n'
            countttter.release()
        else:
            queueLock.release()
       # time.sleep(1)
threadList= []
for fuck in range(0,int(sys.argv[2])):
    threadList.append("Thread-"+str(fuck))
#threadList = ["Thread-1", "Thread-2", "Thread-3"]
#nameList = ["One", "Two", "Three", "Four", "Five"]
queueLock = threading.Lock()
filelock = threading.Lock()
workQueue = Queue.Queue(len(marylines))
threads = []
threadID = 1

# Create new threads
for tName in threadList:
    thread = myThread(threadID, tName, workQueue)
    thread.start()
    threads.append(thread)
    threadID += 1

# Fill the queue
queueLock.acquire()
for word in marylines:
    workQueue.put(word)
queueLock.release()

# Wait for queue to empty
while not workQueue.empty():
    pass

# Notify threads it's time to exit

exitFlag = True

# Wait for all threads to complete
for t in threads:
    t.join()
print "Exiting Main Thread"

关于程序整个流程的解释

首先打开一个文件

进行readline操作,将每一行读取出来,形成一个列表存到marylines变量中

然后根据传入的第二个参数(线程数量),生成以数字进行编号的线程名称(这一部可以省略,因为线程无所谓名称),不进行编号脚本也可以正常工作

然后创建两个锁,一个用来锁队列,一个用来锁文件

创建一个队列,队列长度为marylines列表长度(待处理的域名)

声明一个空的threads列表,用于存储所有的变量,这个在后面join结束进程的时候会用到

在for循环中,启动所有进程同时将进程对象存入threads列表中

给队列加锁,然后往队列里面填充待处理的域名字符串,填充完毕,开锁

在队列空之前一直循环避免程序结束

然后我们来看线程的内容

主函数为process_data

给队列加锁,判断队列空了没

没空则取出来一个域名,然后开锁

打印出当前线程名称以及待处理的域名

用try-catche语句包裹域名解析的代码

使用socket.gethostbyname解析域名,抛出异常则后面的代码不会执行,打印出当前域名无法解析的信息

未抛出异常则给文件加锁,将域名以追加的形式写入到第三个参数指定的文件中

写入完毕,开锁

在try-catche语句外部,

给counter加锁,然后更改counter的值(+1),然后打印出当前counter的值以及总共的值,然后开锁

counter锁在程序最前面进行了声明

最后是判断队列是否为空的else分支:给队列开锁

我试了一下,直接开50个线程,速度还是很快的