mirror of
https://github.com/JHUAPL/CodeCut.git
synced 2026-01-09 13:28:06 -05:00
310 lines
10 KiB
Python
310 lines
10 KiB
Python
##############################################################################################
|
|
# Copyright 2018 The Johns Hopkins University Applied Physics Laboratory LLC
|
|
# All rights reserved.
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy of this
|
|
# software and associated documentation files (the "Software"), to deal in the Software
|
|
# without restriction, including without limitation the rights to use, copy, modify,
|
|
# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
|
|
# permit persons to whom the Software is furnished to do so.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
|
|
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
|
|
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
|
|
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
|
|
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
|
|
# OR OTHER DEALINGS IN THE SOFTWARE.
|
|
#
|
|
# HAVE A NICE DAY.
|
|
|
|
IDA_VERSION = 7
|
|
|
|
if (IDA_VERSION < 7):
|
|
import idc
|
|
import struct
|
|
import idautils
|
|
import basicutils_6x as basicutils
|
|
else:
|
|
import ida_idaapi
|
|
import ida_idc
|
|
import ida_funcs
|
|
import ida_nalt
|
|
import ida_segment
|
|
import idautils
|
|
import basicutils_7x as basicutils
|
|
|
|
import math
|
|
import nltk
|
|
import nltk.collocations
|
|
import re
|
|
|
|
|
|
### NLP Section ###
|
|
|
|
# This section of code attempts to name the modules based on common strings in the string references
|
|
# Not really based on any sound science or anything - your mileage may heavily vary. :-D
|
|
|
|
#string_range_tokenize(start,end,sep):
|
|
#Compile all string references between start and end as a list of strings (called "tokens")
|
|
# <sep> should be a nonsense word, and will show up in the list
|
|
def string_range_tokenize(start,end,sep):
|
|
# get all string references in this range concatenated into a single string
|
|
t = basicutils.CompileTextFromRange(start,end,sep)
|
|
|
|
#Enable this if you already have a bunch of function names and want to include that in the mix
|
|
#t+= basicutils.CompileFuncNamesFromRangeAsText(start,end,sep)
|
|
|
|
#print "string_range_tokenize: raw text:"
|
|
#print t
|
|
#remove printf/sprintf format strings
|
|
tc = re.sub("%[0-9A-Za-z]+"," ",t)
|
|
#convert dash to underscore
|
|
tc = re.sub("-","_",tc)
|
|
#replace _ and / with space - may want to turn this off sometimes
|
|
#this will break up snake case and paths
|
|
#problem is that if you have a path that is used throughout the binary it will probably dominate results
|
|
tc = re.sub("_"," ",tc)
|
|
#replace / and \\ with a space
|
|
tc = re.sub("[/\\\\]"," ",tc)
|
|
#remove anything except alphanumeric, spaces, . (for .c, .cpp, etc) and _
|
|
tc = re.sub("[^A-Za-z0-9_\.\s]"," ",tc)
|
|
|
|
#lowercase it - and store this as the original set of tokens to work with
|
|
tokens = [tk.lower() for tk in tc.split()]
|
|
|
|
#remove English stop words
|
|
#this is the list from the MIT *bow project
|
|
eng_stopw = {"about","all","am","an","and","are","as","at","be","been","but","by","can","cannot","did","do","does","doing","done","for","from","had","has","have","having","if","in","is","it","its","of","on","that","the","these","they","this","those","to","too","want","wants","was","what","which","will","with","would"}
|
|
#remove "code" stop words
|
|
#e.g. common words in debugging strings
|
|
code_sw = {"error","err","errlog","log","return","returned","byte","bytes","status","len","length","size","ok","0x","warning","fail","failed","failure","invalid","illegal","param","parameter","done","complete","assert","assertion","cant","didnt","class","foundation","cdecl","stdcall","thiscall"}
|
|
stopw = eng_stopw.union(code_sw)
|
|
c = 0
|
|
|
|
tokens_f = []
|
|
|
|
for t in tokens:
|
|
if t not in stopw:
|
|
tokens_f.append(t)
|
|
|
|
return tokens_f
|
|
|
|
#bracket_strings(start,end,b_brack,e_brack):
|
|
#Return the most common string in the range <star,end> that begins with b_brack and ends with e_brack
|
|
# The count of how many times this string appeared is also returned
|
|
#I find somewhat often people format debug strings like "[MOD_NAME] Function X did Y!"
|
|
#This function is called by guess_module_names() - if you see this format with different brackets
|
|
#you can edit that call
|
|
def bracket_strings(start,end,b_brack,e_brack):
|
|
sep = "tzvlw"
|
|
t = basicutils.CompileTextFromRange(start,end,sep)
|
|
tokens = [tk.lower() for tk in t.split(sep)]
|
|
|
|
b=[]
|
|
for tk in tokens:
|
|
tk = tk.strip()
|
|
|
|
if tk.startswith(b_brack) :
|
|
b_contents = tk[1:tk.find(e_brack)]
|
|
#Hack to get rid of [-],[+],[*] - could also try to remove non alpha
|
|
if (len(b_contents) > 3):
|
|
#Hack for debug prints that started with [0x%x]
|
|
if (b_contents != "0x%x"):
|
|
b.append(tk[1:tk.find(e_brack)])
|
|
|
|
print("bracket_strings tokens:")
|
|
print(tokens)
|
|
print(b)
|
|
|
|
u_gram=""
|
|
u_gram_score=0
|
|
if (len(b) > 0):
|
|
f = nltk.FreqDist(b)
|
|
u_gram = f.most_common(1)[0][0]
|
|
u_gram_score = f.most_common(1)[0][1]
|
|
|
|
return (u_gram,u_gram_score)
|
|
|
|
#source_file_strings(start,end):
|
|
#Return the most common string that looks like a source file name in the given range
|
|
# The count of how many times this string appeared is also returned
|
|
def source_file_strings(start,end):
|
|
sep = "tzvlw"
|
|
t = basicutils.CompileTextFromRange(start,end,sep)
|
|
#normally would do lower here to normalize but we lose camel case that way
|
|
tokens = [tk for tk in t.split(sep)]
|
|
|
|
#for each string, remove quotes and commas, then tokenize based on spaces to generate the final list
|
|
tokens2=[]
|
|
for tk in tokens:
|
|
tk = tk.strip()
|
|
#strip punctuation, need to leave in _ for filenames and / and \ for paths
|
|
tk = re.sub("[\"\'\,]"," ",tk)
|
|
for tk2 in tk.split(" "):
|
|
tokens2.append(tk2)
|
|
|
|
b=[]
|
|
for tk in tokens2:
|
|
tk = tk.strip()
|
|
if tk.endswith(".c") or tk.endswith(".cpp") or tk.endswith(".cc"):
|
|
#If there's a dir path, only use the end filename
|
|
#This could be tweaked if the directory structure is part of the software architecture
|
|
#e.g. if there are multiple source directories with meaningful names
|
|
if tk.rfind("/") != -1:
|
|
ntk = tk[tk.rfind("/")+1:]
|
|
elif tk.rfind("\\") != -1:
|
|
ntk = tk[tk.rfind("\\")+1:]
|
|
else:
|
|
ntk = tk
|
|
b.append(ntk)
|
|
|
|
print("source_file_strings tokens:")
|
|
#print tokens
|
|
print(b)
|
|
|
|
#a better way to do this (if there are multiple)
|
|
#would be to sort, uniquify, and then make the name foo.c_and_bar.c
|
|
u_gram=""
|
|
u_gram_score=0
|
|
if (len(b) > 0):
|
|
f = nltk.FreqDist(b)
|
|
u_gram = f.most_common(1)[0][0]
|
|
u_gram_score = f.most_common(1)[0][1]
|
|
|
|
return (u_gram,u_gram_score)
|
|
|
|
#common_strings(start,end):
|
|
#Return a list of the common strings in the given range
|
|
#Uses NLTK to generate a list of unigrams, bigrams, and trigrams (1 word, 2 word phrase, 3 word phrase)
|
|
#If the trigram score > 1/2 * bigram score, the most common trigram is used
|
|
#If the bigram score > 1/2 * unigram score, the most common bigram is used
|
|
#Otherwise the most common unigram (single word is used)
|
|
def common_strings(start,end):
|
|
CS_THRESHOLD = 6
|
|
sep = "tvlwz"
|
|
|
|
tokens = string_range_tokenize(start,end,sep)
|
|
|
|
#make a copy since we're going to edit it
|
|
u_tokens = tokens
|
|
c=0
|
|
while (c<len(u_tokens)):
|
|
if u_tokens[c] == sep:
|
|
del u_tokens[c]
|
|
else:
|
|
c+=1
|
|
|
|
print("common_strings tokens:")
|
|
print(tokens)
|
|
|
|
if len(u_tokens) < CS_THRESHOLD:
|
|
#print "%08x - %08x : %s" % (start,end,"no string")
|
|
return ("",0)
|
|
|
|
f = nltk.FreqDist(u_tokens)
|
|
u_gram = f.most_common(1)[0][0]
|
|
u_gram_score = f.most_common(1)[0][1]
|
|
|
|
#print "Tokens:"
|
|
#print tokens
|
|
#print len(tokens)
|
|
|
|
bgs = list(nltk.bigrams(tokens))
|
|
c=0
|
|
while (c<len(bgs)):
|
|
if sep in bgs[c]:
|
|
del bgs[c]
|
|
else:
|
|
c+=1
|
|
|
|
#print "Bigrams:"
|
|
#print bgs
|
|
if (len(bgs) != 0):
|
|
fs = nltk.FreqDist(bgs)
|
|
b_gram = fs.most_common(1)[0][0]
|
|
#print "Most Common:"
|
|
#print b_gram
|
|
b_str = b_gram[0] + "_" + b_gram[1]
|
|
b_gram_score = fs.most_common(1)[0][1]
|
|
else:
|
|
b_str =""
|
|
b_gram_score = 0
|
|
|
|
tgs = list(nltk.trigrams(tokens))
|
|
c=0
|
|
while (c<len(tgs)):
|
|
if sep in tgs[c]:
|
|
del tgs[c]
|
|
else:
|
|
c+=1
|
|
#print "Trigrams:"
|
|
#print tgs
|
|
if (len(tgs) != 0):
|
|
ft = nltk.FreqDist(tgs)
|
|
t_gram = ft.most_common(1)[0][0]
|
|
t_str = t_gram[0] + "_" + t_gram[1] + "_" + t_gram[2]
|
|
t_gram_score = ft.most_common(1)[0][1]
|
|
else:
|
|
t_str = ""
|
|
t_gram_score = 0
|
|
|
|
|
|
#print "1: %s - %d 2: %s - %d 3: %s - %d\n" % (u_gram,u_gram_score,b_str,b_gram_score,t_str,t_gram_score)
|
|
|
|
if (b_gram_score * 2 >= u_gram_score):
|
|
if (t_gram_score * 2 >= b_gram_score):
|
|
ret = t_str
|
|
ret_s = t_gram_score
|
|
else:
|
|
ret = b_str
|
|
ret_s = b_gram_score
|
|
else:
|
|
ret = u_gram
|
|
ret_s = u_gram_score
|
|
|
|
#print "%08x - %08x : %s" % (start,end,ret)
|
|
|
|
return (ret,ret_s)
|
|
|
|
### End of NLP Section ###
|
|
|
|
|
|
|
|
#guess_module_names():
|
|
#Use the NLP section (above) to guess the names of modules and add them to the global module list
|
|
#Attempts to find common bracket strings (e.g. "[MOD_NAME] Debug print!")
|
|
#then source file names (most often left over from calls to assert())
|
|
#then common trigram/bigram/unigrams
|
|
#You can tweak the switchover thresholds below.
|
|
def guess_module_names(module_list):
|
|
#idea - make score threshold based on the size of the module
|
|
# (e.g. smaller modules should have a smaller threshold
|
|
C_SCORE_THRESHOLD = 3
|
|
S_SCORE_THRESHOLD = 1
|
|
B_SCORE_THRESHOLD = 1
|
|
c=0
|
|
unk_mod=0
|
|
while (c<len(module_list)):
|
|
m = module_list[c]
|
|
# first look for strings that start with [FOO], (bracket strings)
|
|
# then look for strings that contain source files (.c,.cpp,etc.)
|
|
# then try common strings
|
|
# above thresholds can be tweaked - they represent the number of strings that have to be repeated
|
|
# in order to use that string as the module name
|
|
(name,scr) = bracket_strings(m.start,m.end,"[","]")
|
|
if (scr < B_SCORE_THRESHOLD):
|
|
(name,scr) = source_file_strings(m.start,m.end)
|
|
if (scr < S_SCORE_THRESHOLD):
|
|
(name,scr) = common_strings(m.start,m.end)
|
|
if (scr < C_SCORE_THRESHOLD):
|
|
#Couldn't come up with a name so name it umod1, umod2, etc.
|
|
name = "umod%d" % (unk_mod)
|
|
#"word cloud" or something to get an idea of what the module is
|
|
#print basicutils.CompileTextFromRange(m.start,m.end," ")
|
|
unk_mod+=1
|
|
module_list[c].name = name
|
|
module_list[c].score = scr
|
|
print("%08x - %08x : %s (%d)" % (m.start,m.end,name,scr))
|
|
c+=1
|
|
|
|
return module_list
|