163 lines
5.1 KiB
Python
163 lines
5.1 KiB
Python
# Sample archive loader: TAR file format
|
|
# Feel free to improve it, this is just a sample
|
|
|
|
import os
|
|
import os.path
|
|
import tarfile
|
|
import shutil
|
|
import tempfile
|
|
import idaapi
|
|
from ida_kernwin import Choose
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
def accept_file(li, filename):
|
|
"""
|
|
Check input file format.
|
|
This function will be called one or more times depending on the result value.
|
|
|
|
@param li: a file-like object which can be used to access the input data
|
|
@param filename: name of the file, if it is an archive member name then the actual file doesn't exist
|
|
@return: 0 - no more supported formats
|
|
string "name" - format name to display in the chooser dialog
|
|
dictionary {
|
|
'format': "name",
|
|
'options': integer,
|
|
'flags': integer
|
|
}
|
|
options: should be 1,
|
|
if ORed with ACCEPT_ARCHIVE then it is an archive loader
|
|
if ORed with ACCEPT_CONTINUE then this function will be called another time
|
|
if ORed with ACCEPT_FIRST then indicates preferred format
|
|
loader_flags: see GENFLG_
|
|
"""
|
|
li.seek(0)
|
|
try:
|
|
t = tarfile.open(fileobj=li, mode='r|*')
|
|
t.close()
|
|
(_, ext) = os.path.splitext(filename)
|
|
if ext not in (".tar", ".tgz", "tar.gz"):
|
|
return 0
|
|
return {'format': "TAR archive",
|
|
'options': 1 | idaapi.ACCEPT_ARCHIVE}
|
|
except Exception:
|
|
pass
|
|
return 0
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
def _read_whole_file(li):
|
|
li.seek(0)
|
|
return li.read(li.size())
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
def _tmpnam():
|
|
(h, n) = tempfile.mkstemp()
|
|
os.close(h)
|
|
return n
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
class TarMemberChoose(Choose):
|
|
"""
|
|
TAR archive members selection chooser
|
|
"""
|
|
|
|
def __init__(self, archive, items):
|
|
title = "Archive: " + archive
|
|
Choose.__init__(
|
|
self,
|
|
title,
|
|
[["File name", Choose.CHCOL_PATH | 60],
|
|
["Size", Choose.CHCOL_DEC | 10]],
|
|
icon=-1, y1=-2,
|
|
flags = Choose.CH_MODAL | Choose.CH_NOIDB)
|
|
self.items = items
|
|
|
|
def OnGetLine(self, n):
|
|
return [self.items[n].name, str(self.items[n].size)]
|
|
|
|
def OnGetSize(self):
|
|
return len(self.items)
|
|
|
|
|
|
# -----------------------------------------------------------------------
|
|
def process_archive(li, archive, defmember, neflags, formatname):
|
|
"""
|
|
Display list of archive members and let the user select one.
|
|
Extract the selected archive member into a temporary file.
|
|
|
|
@param li: a file-like object which can be used to access the input data
|
|
@param archive: name of archive
|
|
@param defmember: extract the specified member
|
|
@param neflags: options selected by the user, see loader.hpp
|
|
@param formatname: name of type of the file
|
|
|
|
@return: '' cancelled by the user
|
|
string error message
|
|
dictionary {
|
|
'temp_file': string,
|
|
'module_name': string,
|
|
'neflags': integer
|
|
}
|
|
temp_file: name of the file with the extracted archive member
|
|
module_name: name of the extracted archive member
|
|
neflags: options selected by the user, see loader.hpp
|
|
"""
|
|
li.seek(0)
|
|
try:
|
|
t = tarfile.open(fileobj=li, mode='r|*')
|
|
except tarfile.TarError as e:
|
|
return str(e)
|
|
|
|
# list of archive members,
|
|
members = t.getmembers()
|
|
t.close()
|
|
|
|
# we are interested in regular files only
|
|
items = [members[i] for i in range(0, len(members))]
|
|
|
|
# if default archive member is specified
|
|
if defmember:
|
|
for m in items:
|
|
if os.path.basename(m.name) == defmember:
|
|
selected_item = m
|
|
break
|
|
else:
|
|
return "Unknown TAR archive default member: %s" % defmember
|
|
else:
|
|
chooser = TarMemberChoose(archive, items)
|
|
code = chooser.Show(True)
|
|
if code == Choose.NO_SELECTION:
|
|
return "" # user canceled
|
|
selected_item = items[code]
|
|
|
|
# construct archive member name
|
|
member_name = os.path.basename(selected_item.name)
|
|
module_name = os.path.join(os.path.dirname(archive), member_name)
|
|
|
|
# file for archive member
|
|
workfile = _tmpnam()
|
|
|
|
# extract member
|
|
# there is a bug reported in 2010 year but not fixed yet:
|
|
# http://bugs.python.org/issue10436
|
|
li.seek(0)
|
|
buf = _read_whole_file(li)
|
|
(h, workfile_tar) = tempfile.mkstemp()
|
|
os.write(h, buf)
|
|
os.close(h)
|
|
|
|
t = tarfile.open(name=workfile_tar, mode='r:*')
|
|
tarinfo = t.getmember(selected_item.name)
|
|
f_in = t.extractfile(tarinfo)
|
|
f_out = open(workfile, 'wb')
|
|
shutil.copyfileobj(f_in, f_out)
|
|
f_out.close()
|
|
f_in.close()
|
|
t.close()
|
|
os.unlink(workfile_tar)
|
|
|
|
return {'temp_file': workfile, 'module_name': module_name}
|