Python安全工具撰寫-pcap流量包重放

最近因為log4j漏洞，作業量急劇增加，卑微小乙沒日沒夜地進行應急處理，為了能測驗是否能檢測到log4j攻擊，我專門寫了一個腳本，當然這個腳本也不只是檢測log4j，只要是pcap包放在指定目錄下則可以重放，

我寫腳本的時候總喜歡先確認功能點，畢竟這是我們寫代碼的前提了，
功能點：將指定目錄中所有pcap檔案全部找出重放
1，找到指定目錄(pcaps)中所有pcap檔案

def find_pcap(dir,filelist):
	for i in os.listdir(dir):#將目錄下所有內容拉出來判斷
		path = os.path.join(dir,i)#拼接單個內容和路徑
		if os.path.isfile(path):#判斷path是否是檔案
			if os.path.splitext(path)[1] == '.pcap':#判斷是否是pcap檔案
				filelist.append(path)#將pcap檔案添加到filelist表中
		elif os.path.isdir(path):#判斷是否是目錄
			newdir = path#拼接成新的目錄路徑
			find_pcap(newdir,filelist)#遞回再次尋找pcap檔案
    return filelist#回傳最侄訓取到的pcap表

2，將pcap檔案全部重放

#這里直接呼叫了系統自帶的tcpreplay工具
def run_pcap(pcap,net_card):
	tmp_pcap = os.path.splitext(os.path.basename(pcap))[0]#提取出檔案名
	print("[+] 正在重放%s" % tmp_pcap)#提示正在重放檔案名
	os.system('tcpreplay -i %s -x 3.0 -K %s ' % (net_card,pcap))#拼接命令net_card是網卡，pcap是檔案

主要功能點已經寫出來了，那么剩下的就是一些優化代碼：
1，加入重放日志，以便排查
2，利用多執行緒的方式，可以同時重放pcap包
3，判斷目錄是否存在pcap檔案
4，將輸出的字體以顏色來顯示高亮
5，指定重放pcaps目錄中的某個或多個目錄中的pcap

代碼

import sys
import os
import os.path
import threading
#匯入sys（獲取輸入的值），os（運行系統命令）.os.path（運行系統檔案操作），threading（多執行緒）模塊

def find_pcap(dir,filelist):
	for i in os.listdir(dir):#將目錄下所有內容拉出來判斷
		path = os.path.join(dir,i)#拼接單個內容和路徑
		if os.path.isfile(path):#判斷path是否是檔案
			if os.path.splitext(path)[1] == '.pcap':#判斷是否是pcap檔案
				filelist.append(path)#將pcap檔案添加到filelist表中
		elif os.path.isdir(path):#判斷是否是目錄
			newdir = path#拼接成新的目錄路徑
			find_pcap(newdir,filelist)#遞回再次尋找pcap檔案
    return filelist#回傳最侄訓取到的pcap表

def run_pcap(pcap,net_card):
	tmp_pcap = os.path.splitext(os.path.basename(pcap))[0]#提取出檔案名
	print("\033[1;32;40m[+] 正在重放%s\033[0m" % tmp_pcap)#提示正在重放檔案名
	os.system('touch %s%s.log' % (log_dir,tmp_pcap))#創建log檔案
	file_log = log_dir + tmp_pcap#拼接log路徑
	os.system('tcpreplay -i %s -x 3.0 -K %s > %s.log 2>&1' % (net_card,pcap,file_log))#拼接命令net_card是網卡，pcap是檔案

def main():
	global current_dir,file_dir,log_dir,filelist#將目錄以及pcap串列都設定為全域變數
	current_dir = os.getcwd()#獲取當前目錄
	file_dir = os.path.join(current_dir,'pcaps/')#獲取pcaps目錄
	log_dir = os.path.join(os.getcwd(),'log/')#獲取日志目錄
	
	filelist = []#初始化pcap表
	choice_file = []#初始化選擇pcap表
if len(sys.argv) >= 2:#判斷用戶是否輸入目錄名稱
	for i in range(1,len(sys.argv)):#將輸入的目錄依次執行以下操作
		path = sys.argv[i]#獲取目錄名稱
		choice_file_dir = os.path.join(file_dir,path)#拼接目錄
		try:
			choice_file = _read_dir_pcap(choice_file_dir, filelist)#嘗試獲取指定目錄下的pcap檔案
		except:
			print("\033[1;31;40m[-] 運行失敗，請檢查是否將目錄放在pcaps目錄中！\033[0m")#如果出現報錯，則檢查是否是輸入有問題或目錄是否不存在
			exit()#退出程式
	else:
		print("\033[1;34m[*] 未選擇目錄，默認重放pcaps目錄下所有pcap包\033[0m")#如果沒有輸入指令，則將pcaps目錄下所有pcap檔案重放
		choice_file = _read_dir_pcap(file_dir,filelist)#獲取所有pcap檔案
		if len(filelist) == 0:#如果pcaps中沒有pcap檔案則退出程式
			print("\033[1;31;40m[-] 運行失敗，請檢查是否將檔案或目錄放在pcaps目錄中！\033[0m")
			exit()
		
		else:
			if len(choice_file) == 0:
				print("\033[1;31;40m[-] 運行失敗，請確認目錄是否有pcap檔案！\033[0m")
				#如果選擇目錄中沒有pcap檔案，則退出程式
				exit()
			
		net_card = input("請輸入重放資料包的網卡名稱：")#輸入從哪個網卡發送流量包
		#設定多執行緒重放
		for i in choice_file:
			t = threading.Thread(target=pcap_poc, args=(i,log_dir,net_card))
			t.start()
		
		print("\033[1;34m[*]如有pcap未重放，可以在log/xxx.log查看重放資訊\033[0m")

if __name__ == "__main__":
	main()

所需目錄

實驗截圖

未指定目錄

查看log檔案