def saveLogBtnClick(self):
logDB = mongoDB.DBConn("shutdown").log
logData = logDB.find({}, {'_id': 0})
logDBCount = logDB.count()
for item in logData:
inputData = str(item).replace("{", "").replace("}", "")
print("inputData : ", inputData)
f = open("log.txt", 'w')
for i in range(1, logDBCount + 1):
data = "NUMBER{} --> {}\n".format(i, inputData)
f.write(data)
f.close()
def checkHashInDB(md5Hash):
md5Hash = md5Hash
checkCount = 0
try:
learn_data = mongoDB.DBConn("shutdown").learn_data
md5InDB = learn_data.find({}, {'hash': 1})
print(md5InDB)
for i in md5InDB:
print(i)
if i == md5Hash:
checkCount = 1
break
if checkCount == 0:
return "NO"
else:
return "YES"
except:
print("[+] Error HERE")
def getFIleMachine(filePath):
opcode_data = mongoDB.DBConn("shutdown").opcode_data
learn_data1 = opcode_data.find({}, {"_id": 0, "hash": 0, "detect": 0})
learn_data2 = []
for item in learn_data1:
print(str(item.keys()))
learn_data2.append(list(item.values()))
# 데이터 셔플하기(섞기)
random.shuffle(learn_data2)
# 학습 전용 데이터와 테스트 전용 데이터 분할하기(2:1 비율)
total_len = len(learn_data2)
print("total_len : ", total_len)
train_len = len(learn_data2)
train_data = []
train_label = []
test_data = []
test_label = []
for i in range(total_len):
data = learn_data2[i][1:10]
data2 = getFileInfo.getFileInfo(
"C:\\Users\\Administrator\\Downloads\\CorsixTHInstaller.exe")
del data2['kind']
del data2['hash']
del data2['detect']
data2 = list(data2.values())
label = learn_data2[i][0]
if i < train_len:
train_data.append(data)
train_label.append(label)
test_data.append(data2)
test_label.append(label)
# 데이터를 학습시키고 예측하기
clf = svm.SVC()
clf.fit(train_data, train_label)
pre = clf.predict(test_data)
print("pre :::::::::::::: ", pre[0])
return pre[0]
def checkHashInDB(md5Hash):
checkCount = 0
try:
opcode_data = mongoDB.DBConn("shutdown").opcode_data
selectDB = opcode_data.find({}, {'_id': 0, 'hash': 1, 'kind': 1})
for i in selectDB:
i = str(i).split(",")
kind = i[0].replace("{'kind': '", "").replace("'", "")
hash = i[1].replace(" 'hash': '", "").replace("'}", "")
if hash == md5Hash:
if kind == "MALWARE":
checkCount = 1
return "YES", "M"
elif kind == "NORMAL":
checkCount = 0
return "YES", "N"
if checkCount == 0:
return "NO", "?"
else:
return "YES"
except:
pass
def selectLog(self):
print("[+] selectLog")
#########################################################################################
# 검사기록 로그를 불러옴
logDB = mongoDB.DBConn("shutdown").log
logDBCount = logDB.count()
logData = logDB.find()
# 불러온 기록 테이블에 보여주기
num = 0
j = 0
for item in logData:
print(item)
item = str(item).split(",")
item[1] = item[1].replace("'date': '", "").replace("'", "")
item[2] = item[2].replace("'filename': '",
"").replace("'",
"").replace("\\\\", "\\")
item[3] = item[3].replace("'result': '", "").replace("'", "")
item[4] = item[4].replace("'state': '", "").replace("'}", "")
showData = []
showData.append(item[1])
showData.append(item[2])
showData.append(item[3])
showData.append(item[4])
print("showData : ", showData)
print("logDBCount : ", logDBCount)
for i in range(0, 4):
print("j : ", j, " i : ", i, "showData : ", showData[i])
self.ui.logTable.setItem(j, i, QTableWidgetItem(showData[i]))
self.ui.logTable.item(j, i).setTextAlignment(Qt.AlignCenter)
j += 1
# 테이블 CSS 설정
self.ui.logTable.horizontalHeader().setStretchLastSection(True)
self.ui.logTable.setEditTriggers(QAbstractItemView.NoEditTriggers)
self.ui.logTable.setColumnWidth(0, 150)
self.ui.logTable.setColumnWidth(1, 300)
self.ui.logTable.setSelectionBehavior(1)
from sklearn import svm, metrics
import PROJ_COD.MongoDB_Connection as mongoDB
import random, re
# 붗꽃의 CSV 데이터 읽어 들이기
csv = []
result = []
opcode_data = mongoDB.DBConn("shutdown").opcode_data
all_data = opcode_data.find({}, {"_id": 0})
# all_data = opcode_data.find()
for line in all_data:
print("line : ", line)
if "MOVDQU" in line.keys():
ababab = str(line).find("MOVDQU")
print(ababab)
print("11111111111111111111111111111111111")
with open("iris.csv", "r", encoding="utf-8") as fp:
# 한 줄씩 읽어 들이기
for line in fp:
# print (line)
line = line.strip() # 줄바꿈 제거
# print(line)
cols = line.split(",") # 쉼표로 자르기
# print(cols)
# 문자열 데이터를 숫자로 변환하기
fn = lambda n: float(n) if re.match(r'^[0-9\.]+$', n) else n
print(fn("10"))
cols = list(map(fn, cols))
print(cols)
csv.append(cols)
def detectionMain(filePath):
# 악성코드 탐지할 경로 가져오기
# pyqt에서 선택한 내용으로 집어 넣기
# Version1: selFile -> 파일 하나 선택한 경우
# Version2: allFiles -> 폴더 경로 선택한 경우
# Version1 : selFile
########################################################################################
# Version2 : allFiles(폴더 선택한 경우)
allFiles = glob.glob(filePath)
# 가져온 리스트만큼 반복작업 진행
for fileName in allFiles:
print("file_name : ", fileName)
hashRslt = filehash.Get_File_Hash_Main(fileName)[0]
hashFile = filehash.Get_File_Hash_Main(fileName)[1]
if hashRslt[0] == "NO":
# print("[+] RESULT IS NO")
# vtotalRslt = conVirustotal.get_mal_kind(fileName)
# print("vtotalRslt : ",vtotalRslt)
# print("vtotalRslt[0] : ",vtotalRslt[0])
# if vtotalRslt[0] > 10:
# print("[+] THIS FILE MIGHT BE A MALWARE. WE HAVE TO CHECK MORE USING OUR MACHINE. PLEASE WAIT. TILL IT'S OVER.")
# # 테스트할 때 필요한 코드(악성 학습데이터 저장하기)
# malInfo.get_info(fileName).update({"kind":"MALWARE","hash":hashFile})
# print("[+] MAL_INFORMATION : ", malInfo.get_info(fileName))
# insert_data = malInfo.get_info(fileName)
# try: del insert_data['_id']
# except: pass
# learn_data = mongoDB.DBConn("shutdown").learn_data
# learn_data.insert(insert_data)
# print("[+] WE FINISHED INSERT MALWARE DATA IN SHUTDOWN.LEARN_DATA SUCCESSFULLY.")
# elif vtotalRslt[0] <= 10:
# try:
# print("[+] THIS FILE IS NOT A MALWARE. PROCESS IS FINISHED.")
# # 테스트할 때 필요한 코드(정상 학습데이터 저장하기)
# malInfo.get_info(fileName).update({"kind": "NORMAL", "hash": hashFile})
# print("[+] MAL_INFORMATION : ", malInfo.get_info(fileName))
# insert_data = malInfo.get_info(fileName)
# try: del insert_data['_id']
# except: pass
# learn_data = mongoDB.DBConn("shutdown").learn_data
# learn_data.insert(insert_data)
# print("[+] WE FINISHED INSERT NORMAL DATA IN SHUTDOWN.LEARN_DATA SUCCESSFULLY.")
# except:
# print("[+] THE ERROR IN MAIN 51")
print("[+] RESULT IS NO")
vtotalRslt = conVirustotal.get_mal_kind(fileName)
print("vtotalRslt : ", vtotalRslt)
print("vtotalRslt[0] : ", vtotalRslt[0])
if vtotalRslt[0] > 10:
print(
"[+] THIS FILE MIGHT BE A MALWARE. WE HAVE TO CHECK MORE USING OUR MACHINE. PLEASE WAIT. TILL IT'S OVER.")
# 테스트할 때 필요한 코드(악성 학습데이터 저장하기)
opInfo.get_info(fileName).update({"kind": "MALWARE", "hash": hashFile})
print("[+] MAL_INFORMATION : ", opInfo.get_info(fileName))
insert_data = opInfo.get_info(fileName)
try:
del insert_data['_id']
except:
pass
opcode_data = mongoDB.DBConn("shutdown").opcode_data
opcode_data.insert(insert_data)
print("[+] WE FINISHED INSERT MALWARE DATA IN SHUTDOWN.LEARN_DATA SUCCESSFULLY.")
elif vtotalRslt[0] <= 10:
try:
print("[+] THIS FILE IS NOT A MALWARE. PROCESS IS FINISHED.")
# 테스트할 때 필요한 코드(정상 학습데이터 저장하기)
opInfo.get_info(fileName).update({"kind": "NORMAL", "hash": hashFile})
print("[+] MAL_INFORMATION : ", opInfo.get_info(fileName))
insert_data = opInfo.get_info(fileName)
try:
del insert_data['_id']
except:
pass
opcode_data = mongoDB.DBConn("shutdown").opcode_data
opcode_data.insert(insert_data)
print("[+] WE FINISHED INSERT NORMAL DATA IN SHUTDOWN.LEARN_DATA SUCCESSFULLY.")
except:
print("[+] THE ERROR IN MAIN 51")
elif hashRslt[0] == "YES":
if hashRslt[1] == "M":
print("[+] THIS FILE IS A MALWARE. WE FIND OUT MATCHED HASH VALUE IN OUR MALWARE_DATABAE")
print("[+] YOU HAVE TO REMOVE THIS FILE")
elif hashRslt[1] == "N":
print("[+] THIS FILE IS NOT A MALWARE. PROCESS IS FINISHED.")
else:
print("[+] GET_FILE_HASH_MAIN ERROR")
# break
# rslt = vs.get_mal_kind(i)
# print(rslt)
# learnDB = mg.DBConn("shutdown").learnDB
# if rslt[0] > 0:
# if fh.get_hash_match(i) == False:
# insert_test_doc = get_info()
# try:
# del insert_test_doc['_id']
# except:
# pass
# insert_test_doc.update({"detect": rslt[1], "hash": rslt[2]})
# print("insert_teset_doc : ", insert_test_doc)
# try:
# users.insert(insert_test_doc)
# print("[+] insert success", sys.exc_info()[0])
# except:
# print("[-] insert failed", sys.exc_info()[0])
# else:
# print("[-] We find matched Hash. You don't need to insert again!")
import pymongo
import tensorflow as tf
import numpy as np
import os
import time
import datetime
import PROJ_COD.Machine.lee.data_helpers as data_helpers
from PROJ_COD.Machine.lee.text_cnn import TextCNN
from tensorflow.contrib import learn
import PROJ_COD.MongoDB_Connection as mg
# 몽고디비 연결 =======================================================
db = mg.DBConn("shutdownDB")
learning_mal_data = db.mal_info #학습시킬 악성코드 정보 디비 연결
learning_nml_data = db.nml_info #학습시킬 정상파일 정보 디비 연결
# =====================================================================
# Parameters (기본값)==================================================
# Model Hyperparameters
tf.flags.DEFINE_integer(
"embedding_dim", 128,
"Dimensionality of character embedding (default: 128)")
tf.flags.DEFINE_string("filter_sizes", "2,3,4",
"Comma-separated filter sizes (default: '3,4,5')")
tf.flags.DEFINE_integer("num_filters", 128,
"Number of filters per filter size (default: 128)")
tf.flags.DEFINE_float("dropout_keep_prob", 0.5,
"Dropout keep probability (default: 0.5)")
tf.flags.DEFINE_float("l2_reg_lambda", 0.0,
"L2 regularizaion lambda (default: 0.0)")
# Training parameters
def startFileDetectionBtnClick(self):
todayTime = datetime.today().strftime("%Y-%m-%d %H:%M:%S")
self.ui.page2_dateLabel.setText(todayTime)
fileBaseName = os.path.basename(fname)
self.ui.page2_filenameLabel.setText(fileBaseName)
fileDirName = os.path.dirname(fname)
self.ui.page2_filepathLabel.setText(fileDirName)
fileHash = getFileHash.getFileHash(fname)
self.ui.page2_md5Label.setText(fileHash)
#사용자가 선택한 파일의 Hash 값과 DB에 있는 Hash값을 비교하여 등록 여부 및 악성코드 여부 확인
md5Check = getFileHash.checkHashInDB(fname)
if md5Check[0] == "YES":
if md5Check[1] == "M":
self.ui.MD5HashLabel.setText("악성 코드")
self.ui.page2_similarLabel.setText("100%")
self.ui.numberCountLabel.setPixmap(QPixmap("그림1.jpg"))
elif md5Check[1] == "N":
self.ui.MD5HashLabel.setText("정상 파일")
self.ui.page2_similarLabel.setText("0%")
self.ui.numberCountLabel.setPixmap(QPixmap("그림1.jpg"))
else:
pass
elif md5Check[0] == "NO":
self.ui.page2_similarLabel.setText("100%")
self.ui.MD5HashLabel.setText("미등록")
self.ui.numberCountLabel.setPixmap(QPixmap("그림1.jpg"))
# 머신러닝 돌리는 코드
# 파일 경로 전송해줘야됨
machineRslt = fileMachine.getFIleMachine(fname)
print(
"1111111111111111111111111111111111111111111111111111111111111 ",
machineRslt)
if machineRslt == "NOMAL":
self.ui.page2_similarLabel.setText("정상파일")
self.ui.machineResultLabel.setText("정상파일")
else:
self.ui.page2_similarLabel.setText("악성파일")
self.ui.machineResultLabel.setText("악성파일")
self.ui.numberCountLabel.setPixmap(QPixmap("그림3.jpg"))
################################
# main = vt.Virustotal()
# dic = main.rscReport(md5Check)
# try:
# if dic['positives'] > 10:
# for key, value in dic['scans'].items():
# if value['result'] != None:
# ss=dic['positives'], value['result'], dic['md5']
# self.ui.numCountLabel2.setText("악성 코드")
# self.ui.numberCountLabel.setPixmap(QPixmap("그림2.jpg"))
#
# else:
# self.ui.numCountLabel2.setText("정상 파일")
# self.ui.numberCountLabel.setPixmap(QPixmap("그림2.jpg"))
# except:
# self.ui.numCountLabel2.setText("Error")
# self.ui.numberCountLabel.setPixmap(QPixmap("그림2.jpg"))
# print("[+] dic['positives'] DOES NOT EXIST.")
# ################################
# rrr= vt.get_mal_kind(fname)
# print(rrr)
# else:
# print("Error !!!")
#정상 파일일 경우 삭제버튼 비활성화
if self.ui.page2_similarLabel.text() == "0%":
self.ui.delFileBtn.setEnabled(False)
self.ui.delFileBtn.setStyleSheet(
"background-color:lightgray;color: white;border:nono;")
#악성 파일일 경우 삭제버튼 활성화
else:
self.ui.delFileBtn.setEnabled(True)
self.ui.delFileBtn.setStyleSheet(
"background-color: #0F75BD;color: white;border:nono;")
#########################################################################################
Form.selectLog(self)
# 검사시작 버튼 클릭시 tab3에 보여줄 로그 데이터베이스에 INSERT
logDB = mongoDB.DBConn("shutdown").log
if logDB.count() == 10:
logDB.remove()
insertLogData = {}
insertLogData.update({
"date": todayTime,
"filename": fname,
"result": "normal",
"state": "finished"
})
logDB.insert(insertLogData)
def delLogBtnClick(self):
logDB = mongoDB.DBConn("shutdown").log
logDB.remove()
for j in range(0, 11):
for i in range(0, 4):
self.ui.logTable.setItem(j, i, QTableWidgetItem(""))
def getFileInfo(filePath):
pe = pefile.PE(filePath)
op_list_count = {"kind": "", "hash": "","detect":"", "BT": 0, "FDIVP": 0, "FILD": 0, "IMUL": 0, "INT": 0, "NOP": 0, "SBB": 0, "SHLD": 0, "STD": 0}
op_list_count_md = {"kind": "", "hash": "","detect":"", "BT": 0, "FDIVP": 0, "FILD": 0, "IMUL": 0, "INT": 0, "NOP": 0, "SBB": 0, "SHLD": 0, "STD": 0}
for section in pe.sections:
flags = []
for flag in sorted(section_flags):
if getattr(section, flag[0]):
flags.append(flag[0])
if 'IMAGE_SCN_MEM_EXECUTE' in flags:
iterable = distorm3.DecodeGenerator(0, section.get_data(), distorm3.Decode32Bits)
for (offset, size, instruction, hexdump) in iterable:
op_code = instruction.split()[0]
op_code = str(op_code).lstrip('b')
op_code = str(op_code).replace("'", "")
if op_code not in op_list_count.keys():
op_list_count[op_code] = 1
elif op_code in op_list_count.keys():
op_list_count[op_code] = op_list_count[op_code] + 1
for flag in sorted(section_flags):
if getattr(section, flag[0]):
flags.append(flag[0])
s_name1 = str(section.Name)
s_name = re.sub(r"[b'|\\x00]", "", s_name1)
if s_name == '.tet':
s_name = '.text'
s_name = s_name.replace(".", "_")
pe.parse_data_directories(
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'])
print(op_list_count)
op_list_count_md.update({
"kind": op_list_count.get("kind"),
"hash": op_list_count.get("hash"),
"detect": op_list_count.get("detect"),
"BT": op_list_count.get("BT"),
"FDIVP": op_list_count.get("FDIVP"),
"FILD": op_list_count.get("FILD"),
"IMUL": op_list_count.get("IMUL"),
"INT": op_list_count.get("INT"),
"NOP": op_list_count.get("NOP"),
"SBB": op_list_count.get("SBB"),
"SHLD": op_list_count.get("SHLD"),
"STD": op_list_count.get("STD"),
})
print(op_list_count_md)
# 가져온 리스트만큼 반복작업 진행
rslt = vs.get_mal_kind(filePath)
print(rslt)
rslt2 = gfh.checkHashInDB(rslt[2])
print(rslt2)
if rslt2[0] == "NO":
opcode_data = mg.DBConn("shutdown").opcode_data
try:
del op_list_count_md['_id']
except:pass
if rslt[0]>9:
op_list_count_md.update({"kind":"MALWARE","detect": rslt[1], "hash" : rslt[2]})
else:
op_list_count_md.update({"kind": "NORMAL", "detect": "NORMAL", "hash": rslt[2]})
print("insert_teset_doc : ",op_list_count_md)
try:
opcode_data.insert(op_list_count_md)
print("[+] insert success", sys.exc_info()[0])
except:
print("[-] insert failed",sys.exc_info()[0])
return op_list_count_md
"ie_api": str(api_list)
})
return insert_test_doc
if __name__ == "__main__":
# DB에 저장한 샘플 악성코드 폴더의 악성코드 리스트를 가져옴
filelist = glob.glob('C:\\TMP2\\*.exe')
# print(filelist)
# 가져온 리스트만큼 반복작업 진행
for i in filelist:
print("file_name : ", i)
rslt = vs.get_mal_kind(i)
print(rslt)
users = mg.DBConn("maldb").users
if rslt[0] > 0:
if fh.get_hash_match(i) == False:
insert_test_doc = get_info()
try:
del insert_test_doc['_id']
except:
pass
insert_test_doc.update({"detect": rslt[1], "hash": rslt[2]})
print("insert_teset_doc : ", insert_test_doc)
try:
users.insert(insert_test_doc)
print("[+] insert success", sys.exc_info()[0])
except:
print("[-] insert failed", sys.exc_info()[0])
else:
import hashlib, glob
import PROJ_COD.MongoDB_Connection as mongoDB
"""
작성일 : 2017-07-14(최초작성)
수정일 : 2017-07-18(디비연결 공통으로 변경)
버전 : v1.1(수정시 버전 올려주세요)
프로그램 : 파일 해시값을 중복되지 않게 몽고디비에 저장하는 코드
"""
users = mongoDB.DBConn("maldb").users
def getFileHash(filename):
insertFormat = {}
myMd5 = hashlib.md5()
a = 0
with open(filename, 'rb') as checkFile:
for chunk in iter(lambda: checkFile.read(8192), ''):
myMd5.update(chunk)
fileHash = myMd5.hexdigest()
insertFormat = {'hash': fileHash}
item = users.find()
for i in item:
if (i['hash'] == fileHash):
print("[+] find match data : ", i['hash'], fileHash)
a = 1
break
else:
print("[+] not match : ", i['hash'], fileHash)
if a == 0:
try: