pickle.py

import os
import sys
import datetime
import logging
import shutil
import cPickle

def main():

	if len(sys.argv) < 4:

		print "Too few arguments!"
		print "Usage: ", sys.argv[0], "<apk file directory> <temp directory> <opseq directory> <include support libraries>"
		return

	# Location of apk files that need extracting
	apk_file_directory = sys.argv[1]
	print "Reading apks from : ", apk_file_directory

	# Temporary folder for storing extracted apk
	tmp_file_directory = sys.argv[2]
	print "Temp directory for intermediate files : ", tmp_file_directory

	# Output folder for storing the opseqs
	opseq_file_directory = sys.argv[3]
	print "Output for storing opseqs : ", opseq_file_directory

	# Default is not to include smali files in android support libraries unless 4th parameter is provided
	include_support = False
	if len(sys.argv) == 5:
		include_support = ((sys.argv[4]) == "incl")
		print "Android support library smali files will be included"
	else:
		print "Android support library smali files will not be included"

	# A log file is created in the temporary directory
	logging.basicConfig(filename = tmp_file_directory+'/opseq.log', level=logging.DEBUG)
	
	# List storing names of all apks
	apks = []

	# Looping through all apks in apk_directory
	for file in os.listdir(apk_file_directory):
		if os.path.isfile(os.path.join(apk_file_directory, file)):
			apks.append(file)

	logging.info('Total apks to be decoded {0}'.format(len(apks)))
	print "Total apks to be decoded : ",len(apks)

	num = 0
	before = datetime.datetime.now()
	logging.info('Starting at: {0}'.format(before))
	print 'Starting at: {0}',before
	max_length = 0
	# Looping through all apks
	for apk_hash in apks:
		apk_file_location = os.path.join(apk_file_directory, apk_hash)
		num += 1
		logging.info('Decoding apk: {0} apk #: {1}'.format(apk_file_location,num))
		print "apk #: ", num
		print "apk location: ", apk_file_location

		temp_location = None
		# Extracing apk into the temporary directory
		temp_location = extract_apk(apk_file_location, tmp_file_directory,apk_hash,include_support)

		if (not os.path.exists(temp_location) or not os.listdir(temp_location)):
			print "Could not find Smali directory..."
			logging.error('NOT decoded directory: {0}'.format(apk_file_location))
			print "NOT decoded directory: ", apk_file_location
			continue

		# Generating opcode sequence for the apks
		logging.info('Extracting opcode sequence from apk: {0}'.format(apk_hash))
		result, length = generate_opcode_seq(temp_location, opseq_file_directory,apk_hash)
		if max_length < length:
			max_length = length

		if result:
			print "Opseq file for apk #",num," is created"
			logging.info('opseq file for apk# {0} is created'.format(num))
		else:
			logging.error('opseq file creation was unsuccessful')
			print "Opseq file creation was unsuccessful"

		# Deleted extracted apk after extracting opseqs
		if os.path.exists(temp_location):
			shutil.rmtree(temp_location)

	print max_length
	
	after = datetime.datetime.now()
	print "Finished by: {0} ",after
	logging.info('Total time taken: {0}'.format(after-before))
	print "Total time taken: ", after-before

def generate_opcode_seq(temp_dir, opseq_file_directory, apk_hash):
	# Returns True if creating opcode sequence file was successful,
	# searches all files in smali folder,
	# writes the corresponding opcode sequencces into a .opseq file
	# and removes the support library files based on the last parameter 'incl'
	dalvik_opcodes = {}
	# Reading the Dalvik opcodes into a dictionary
	with open("DalvikOpcodes.txt") as file:
		for line in file:
			(key, val) = line.split()
			dalvik_opcodes[key] = val
	try:
		smali_dir = os.path.join(temp_dir, "smali")
		# opseq_fname = os.path.join(opseq_file_directory,apk_hash+".opseq")

		# with open(opseq_fname,"a") as opseq_file:
			# Looping through all smali files
		x = []
		for root, dirs, fnames in os.walk(smali_dir):
			for fname in fnames:
				full_path = os.path.join(root, fname)
				logging.info('Smali file path: {0}'.format(full_path))
				x += get_opcode_seq(full_path, dalvik_opcodes)
				# Generating opcode sequence for the smali file
				# opseq_file.write(get_opcode_seq(full_path, dalvik_opcodes))
		# opseq_file.close()
		
		data_point = {}
		data_point['x'] = [0]*100000
		if(len(x)<100000):
			data_point['x'][:len(x)] = x
		else:
			data_point['x'] = x[:100000]
		print len(data_point['x'])
		data_point['y'] = 0
	
		fp = open(os.path.join(opseq_file_directory,apk_hash+".pickle"),'wb')
		cPickle.dump(data_point, fp, protocol = cPickle.HIGHEST_PROTOCOL)

		return True, len(x)
	except Exception as e:
		print "Exception occured during opseq creation of apk : ", apk_hash
		logging.error('Exception occured during opseq creation {0}'.format(str(e)))
		return False

def get_opcode_seq(smali_fname, dalvik_opcodes):
	# Returns opcode sequence generated from the smali file 'smali_fname'

	# opcode_seq = ''
	x_ = []
	method_cnt = 0
	opcode_cnt = 0
	opcode_global_cnt = 0

	with open(smali_fname, mode="r") as file:
		reader = file.read()
		for i, part in enumerate(reader.split(".method")):
			# add_newline = False
			if i!=0:
				method_cnt = method_cnt + 1
				logging.info('Method No: {0}'.format(str(method_cnt)))
				method_part = part.split(".end method")[0]
				method_body = method_part.strip().split('\n')
				for line in method_body:
					# Ignore everything other than the opcodes
					if not line.strip().startswith('.') and not line.strip().startswith('#') and line.strip():
						method_line = line.strip().split()
						if method_line[0] in dalvik_opcodes:
							opcode_cnt = opcode_cnt + 1
							logging.info('Opcode instruction: {0}'.format(str(method_line[0])))
							logging.info('Integer Equivalent: {0}'.format(str(int(dalvik_opcodes[method_line[0]],16))))
							# add_newline = True
							x_.append(int(dalvik_opcodes[method_line[0]],16))
				logging.info('Total number of opcodes in method# {0} is {1}'.format(str(method_cnt),str(opcode_cnt)))
				opcode_global_cnt = opcode_global_cnt + opcode_cnt
				# Add newline after every method
				# if add_newline:
				# 	opcode_seq += '\n'
	# return opcode_seq
	logging.info('Total number of opcodes in this apk: {0}'.format(str(opcode_global_cnt)))
	return x_

def extract_apk (apk_file_location,tmp_file_directory,hash,include):
	# Extracts the apk at apk_file_location and
	# stores the decoded folders in tmp_file_directory

	output_file_location = os.path.join(tmp_file_directory, hash+ ".smali")
	try:
		apktool_decode_app(apk_file_location,output_file_location,include)
	except ApkToolException:
		print "ApkToolException on decoding"
		logging.error("ApkToolException on decoding apk {0} ".format(apk_file_location))
		pass
	return output_file_location

def apktool_decode_app(apk_file, output_file,include):
	# Runs the apktool on a given apk

	apktooldir = "/usr/local/bin"

	apktoolcmd = "{0}/apktool d -f {1} -o {2}".format(apktooldir, apk_file, output_file)
	result = os.system(apktoolcmd)
	if result!=0 :
		raise ApkToolException(apktoolcmd)

	# Checks if smali files belonging to the android support libraries must be ignored
	if not include:
		# Don't keep the smali files from sndroid support libraries
		android_support_directory = os.path.join(output_file, "smali/android")
		if os.path.exists(android_support_directory):
			rm_cmd = "rm -r %s" %(android_support_directory)
			os.system(rm_cmd)

# Exception class to represent an Apktool Exception
class ApkToolException(Exception):
	def __init__(self, command):
		self.command = command

	def __str__(self):
		return repr(self.command)

if __name__ == '__main__':
	main()