I am trying to process one dataset of files (APKs) as present in one of my directory input/ and for output, I can get the final destination in a directory named output/.
Here presents the python file code as process_dataset.py. In this file, I am extracting the API features and User features from an APK.
import argparse
import json
import multiprocessing
import os
import sys
import traceback
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Union, Optional
import dgl
import joblib as J
import networkx as nx
import torch
from androguard.core.analysis.analysis import MethodAnalysis
from androguard.core.api_specific_resources import load_permission_mappings
from androguard.misc import AnalyzeAPK
from pygtrie import StringTrie
ATTRIBUTES = ['external', 'entrypoint', 'native', 'public', 'static', 'codesize', 'api', 'user']
package_directory = os.path.dirname(os.path.abspath(__file__))
stats: Dict[str, int] = defaultdict(int)
def memoize(function):
"""
Alternative to @lru_cache which could not be pickled in ray
: param function: Function to be cached
: return: Wrapped function
"""
memo = {}
def wrapper(*args):
if args in memo:
return memo[args]
else:
rv = function(*args)
memo[args] = rv
return rv
return wrapper
class FeatureExtractors:
NUM_PERMISSION_GROUPS = 20
NUM_API_PACKAGES = 226
NUM_OPCODE_MAPPINGS = 21
@staticmethod
def _get_opcode_mapping() -> Dict[str, int]:
"""
Group opcodes and assign them an ID
:return: Mapping from opcode group name to their ID
"""
mapping = {x: i for i, x in enumerate(['nop', 'mov', 'return',
'const', 'monitor', 'check-cast', 'instanceof', 'new',
'fill', 'throw', 'goto/switch', 'cmp', 'if', 'unused',
'arrayop', 'instanceop', 'staticop', 'invoke',
'unaryop', 'binop', 'inline'])}
mapping['invalid'] = -1
return mapping
@staticmethod
@memoize
def _get_instruction_type(op_value: int) -> str:
"""
Get instruction group name from instruction
:param op_value: Opcode value
:return: String containing ID of :instr:
"""
if 0x00 == op_value:
return 'nop'
elif 0x01 <= op_value <= 0x0D:
return 'mov'
elif 0x0E <= op_value <= 0x11:
return 'return'
elif 0x12 <= op_value <= 0x1C:
return 'const'
elif 0x1D <= op_value <= 0x1E:
return 'monitor'
elif 0x1F == op_value:
return 'check-cast'
elif 0x20 == op_value:
return 'instanceof'
elif 0x22 <= op_value <= 0x23:
return 'new'
elif 0x24 <= op_value <= 0x26:
return 'fill'
elif 0x27 == op_value:
return 'throw'
elif 0x28 <= op_value <= 0x2C:
return 'goto/switch'
elif 0x2D <= op_value <= 0x31:
return 'cmp'
elif 0x32 <= op_value <= 0x3D:
return 'if'
elif (0x3E <= op_value <= 0x43) or (op_value == 0x73) or (0x79 <= op_value <= 0x7A) or (
0xE3 <= op_value <= 0xED):
return 'unused'
elif (0x44 <= op_value <= 0x51) or (op_value == 0x21):
return 'arrayop'
elif (0x52 <= op_value <= 0x5F) or (0xF2 <= op_value <= 0xF7):
return 'instanceop'
elif 0x60 <= op_value <= 0x6D:
return 'staticop'
elif (0x6E <= op_value <= 0x72) or (0x74 <= op_value <= 0x78) or (0xF0 == op_value) or (
0xF8 <= op_value <= 0xFB):
return 'invoke'
elif 0x7B <= op_value <= 0x8F:
return 'unaryop'
elif 0x90 <= op_value <= 0xE2:
return 'binop'
elif 0xEE == op_value:
return 'inline'
else:
return 'invalid'
@staticmethod
def _mapping_to_bitstring(mapping: List[int], max_len) -> torch.Tensor:
"""
Convert opcode mappings to bitstring
:param max_len:
:param mapping: List of IDs of opcode groups (present in an method)
:return: Binary tensor of length `len(opcode_mapping)` with value 1 at positions specified by :poram mapping:
"""
size = torch.Size([1, max_len])
if len(mapping) > 0:
indices = torch.LongTensor([[0, x] for x in mapping]).t()
values = torch.LongTensor([1] * len(mapping))
tensor = torch.sparse.LongTensor(indices, values, size)
else:
tensor = torch.sparse.LongTensor(size)
# Sparse tensor is normal tensor on CPU!
return tensor.to_dense().squeeze()
@staticmethod
def _get_api_trie() -> StringTrie:
apis = open(Path(package_directory).parent / "metadata" / "api.list").readlines()
api_list = {x.strip(): i for i, x in enumerate(apis)}
api_trie = StringTrie(separator='.')
for k, v in api_list.items():
api_trie[k] = v
return api_trie
@staticmethod
@memoize
def get_api_features(api: MethodAnalysis) -> Optional[torch.Tensor]:
if not api.is_external():
return None
api_trie = FeatureExtractors._get_api_trie()
name = str(api.class_name)[1:-1].replace('/', '.')
_, index = api_trie.longest_prefix(name)
if index is None:
indices = []
else:
indices = [index]
feature_vector = FeatureExtractors._mapping_to_bitstring(indices, FeatureExtractors.NUM_API_PACKAGES)
return feature_vector
@staticmethod
@memoize
def get_user_features(user: MethodAnalysis) -> Optional[torch.Tensor]:
if user.is_external():
return None
opcode_mapping = FeatureExtractors._get_opcode_mapping()
opcode_groups = set()
for instr in user.get_method().get_instructions():
instruction_type = FeatureExtractors._get_instruction_type(instr.get_op_value())
instruction_id = opcode_mapping[instruction_type]
if instruction_id >= 0:
opcode_groups.add(instruction_id)
# 1 subtraction for 'invalid' opcode group
feature_vector = FeatureExtractors._mapping_to_bitstring(list(opcode_groups), len(opcode_mapping) - 1)
return torch.LongTensor(feature_vector)
def process(source_file: Path, dest_dir: Path):
try:
file_name = source_file.stem
_, _, dx = AnalyzeAPK(source_file)
cg = dx.get_call_graph()
mappings = {}
for node in cg.nodes():
features = {
"api": torch.zeros(FeatureExtractors.NUM_API_PACKAGES),
"user": torch.zeros(FeatureExtractors.NUM_OPCODE_MAPPINGS)
}
if node.is_external():
features["api"] = FeatureExtractors.get_api_features(node)
else:
features["user"] = FeatureExtractors.get_user_features(node)
mappings[node] = features
nx.set_node_attributes(cg, mappings)
cg = nx.convert_node_labels_to_integers(cg)
dg = dgl.from_networkx(cg, node_attrs=ATTRIBUTES)
dest_dir = dest_dir / f'{file_name}.fcg'
dgl.data.utils.save_graphs(str(dest_dir), [dg])
print(f"Processed {source_file}")
except:
print(f"Error while processing {source_file}")
traceback.print_exception(*sys.exc_info())
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Preprocess APK Dataset into Graphs')
parser.add_argument(
'-s', '--source-dir',
help='The directory containing apks',
required=True
)
parser.add_argument(
'-d', '--dest-dir',
help='The directory to store processed graphs',
required=True
)
parser.add_argument(
'--override',
help='Override existing processed files',
action='store_true'
)
parser.add_argument(
'--dry',
help='Run without actual processing',
action='store_true'
)
parser.add_argument(
'--n-jobs',
default=multiprocessing.cpu_count(),
help='Number of jobs to be used for processing'
)
parser.add_argument(
'--limit',
help='Run for n apks',
default=-1
)
args = parser.parse_args()
source_dir = Path(args.source_dir)
if not source_dir.exists():
raise FileNotFoundError(f'{source_dir} not found')
dest_dir = Path(args.dest_dir)
if not dest_dir.exists():
raise FileNotFoundError(f'{dest_dir} not found')
n_jobs = args.n_jobs
if n_jobs < 2:
print(f"n_jobs={n_jobs} is too less. Switching to number of CPUs in this machine instead")
n_jobs = multiprocessing.cpu_count()
files = [x for x in source_dir.iterdir() if x.is_file()]
source_files = set([x.stem for x in files])
dest_files = set([x.name for x in dest_dir.iterdir() if x.is_file()])
unprocessed = [source_dir / f'{x}.apk' for x in source_files - dest_files]
print(f"Only {len(unprocessed)} out of {len(source_files)} remain to be processed")
if args.override:
print(f"--override specified. Ignoring {len(source_files) - len(unprocessed)} processed files")
unprocessed = [source_dir / f'{x}.apk' for x in source_files]
print(f"Starting dataset processing with {n_jobs} Jobs")
limit = int(args.limit)
if limit != -1:
print(f"Limiting dataset processing to {limit} apks.")
unprocessed = unprocessed[:limit]
if not args.dry:
J.Parallel(n_jobs=n_jobs)(J.delayed(process)(x, dest_dir) for x in unprocessed)
print("DONE")
After running this code like given below,
python process_dataset.py -s input/ -d output/
I am getting an error as
AttributeError: 'EncodedMethod' object has no attribute 'is_external'
Please, anyone, help with it. I tried to change the networkx version from 2.x to 1.x also but it doesn't change the error. Also, I tried
pip uninstall networkx
pip install dgl
pip install networkx==1.11
