Optimize code for gpt pltools Port utils.py

This commit is contained in:
ColdWindScholar 2024-05-04 14:30:18 +08:00
parent ca140c5dfe
commit 238368d945
4 changed files with 21 additions and 30 deletions

View file

@ -116,7 +116,7 @@ class Port(metaclass=LogBase):
sys.stdout.flush() sys.stdout.flush()
except Exception as serr: except Exception as serr:
print("Handshake: " + str(serr)) print(f"Handshake: {str(serr)}")
if "access denied" in str(serr): if "access denied" in str(serr):
self.warning(str(serr)) self.warning(str(serr))
self.debug(str(serr)) self.debug(str(serr))
@ -194,13 +194,13 @@ class Port(metaclass=LogBase):
else: else:
cmdrsp = self.usbread(dlen) cmdrsp = self.usbread(dlen)
if cmdrsp[0] is not value[0]: if cmdrsp[0] is not value[0]:
self.error("Cmd error :" + hexlify(cmdrsp).decode('utf-8')) self.error(f"Cmd error :{hexlify(cmdrsp).decode('utf-8')}")
return -1 return -1
if bytestoread > 0: if bytestoread > 0:
resp = self.usbread(bytestoread) resp = self.usbread(bytestoread)
return resp return resp
else: else:
self.warning("Couldn't send :" + hexlify(value).decode('utf-8')) self.warning(f"Couldn't send :{hexlify(value).decode('utf-8')}")
return resp return resp
def echo(self, data): def echo(self, data):

View file

@ -328,9 +328,8 @@ class gpt(metaclass=LogBase):
def print_gptfile(self, filename): def print_gptfile(self, filename):
try: try:
filesize = os.stat(filename).st_size
with open(filename, "rb") as rf: with open(filename, "rb") as rf:
size = min(32 * 4096, filesize) size = min(32 * 4096, os.stat(filename).st_size)
data = rf.read(size) data = rf.read(size)
for sectorsize in [512, 4096]: for sectorsize in [512, 4096]:
result = self.parse(data, sectorsize) result = self.parse(data, sectorsize)

View file

@ -76,7 +76,7 @@ class PLTools(metaclass=LogBase):
ack = self.exploit.runpayload(payload, ack, addr, dontack) ack = self.exploit.runpayload(payload, ack, addr, dontack)
if ack == ack: if ack == ack:
self.info("Successfully sent payload: " + filename) self.info(f"Successfully sent payload: {filename}")
self.mtk.daloader.patch = True self.mtk.daloader.patch = True
return True return True
elif ack == b"\xc1\xc2\xc3\xc4": elif ack == b"\xc1\xc2\xc3\xc4":
@ -87,20 +87,20 @@ class PLTools(metaclass=LogBase):
print_progress(0, 100, prefix='Progress:', suffix='Complete', bar_length=50) print_progress(0, 100, prefix='Progress:', suffix='Complete', bar_length=50)
for pos in range(0, 0x40000, 64): for pos in range(0, 0x40000, 64):
wf.write(self.mtk.port.usbread(64)) wf.write(self.mtk.port.usbread(64))
self.info("Preloader dumped as: " + "preloader.bin") self.info("Preloader dumped as: preloader.bin")
return True return True
else: else:
with open("out.bin", 'wb') as wf: with open("out.bin", 'wb') as wf:
print_progress(0, 100, prefix='Progress:', suffix='Complete', bar_length=50) print_progress(0, 100, prefix='Progress:', suffix='Complete', bar_length=50)
for pos in range(0, 0x20000, 64): for pos in range(0, 0x20000, 64):
wf.write(self.mtk.port.usbread(64)) wf.write(self.mtk.port.usbread(64))
self.info("Bootrom dumped as: " + "out.bin") self.info("Bootrom dumped as: out.bin")
return True return True
self.error("Error on sending payload: " + filename) self.error(f"Error on sending payload: {filename}")
return False return False
else: else:
self.error("Error on sending payload: " + filename) self.error(f"Error on sending payload: {filename}")
self.error("Error, payload answered instead: " + hexlify(ack).decode('utf-8')) self.error(f"Error, payload answered instead: {hexlify(ack).decode('utf-8')}")
return False return False
def runbrute(self, args): def runbrute(self, args):
@ -138,16 +138,16 @@ class PLTools(metaclass=LogBase):
self.info("Kamakiri / DA Run") self.info("Kamakiri / DA Run")
if self.runpayload(filename=pfilename, ack=0xC1C2C3C4, offset=0): if self.runpayload(filename=pfilename, ack=0xC1C2C3C4, offset=0):
if self.exploit.dump_brom(filename): if self.exploit.dump_brom(filename):
self.info("Dumped as: " + filename) self.info(f"Dumped as:{filename} ")
return True return True
else: else:
self.error("Error on sending payload: " + filename) self.error(f"Error on sending payload: {filename}")
else: else:
if self.exploit.dump_brom(filename, length=length): if self.exploit.dump_brom(filename, length=length):
self.info("Dumped as: " + filename) self.info(f"Dumped as: {filename}")
return True return True
else: else:
self.error("Error on sending payload: " + pfilename) self.error(f"Error on sending payload: {pfilename}")
return False return False
def run_dump_preloader(self, filename): def run_dump_preloader(self, filename):
@ -158,11 +158,11 @@ class PLTools(metaclass=LogBase):
data, filename = self.exploit.dump_preloader() data, filename = self.exploit.dump_preloader()
return data, filename return data, filename
else: else:
self.error("Error on sending payload: " + pfilename) self.error(f"Error on sending payload: {pfilename}")
return None, None return None, None
else: else:
if self.exploit.dump_brom(filename): if self.exploit.dump_brom(filename):
self.info("Preloader dumped as: " + filename) self.info(f"Preloader dumped as: {filename}")
return True return True
else: else:
self.error("Error on dumping preloader") self.error("Error on dumping preloader")

View file

@ -323,17 +323,13 @@ def do_tcp_server(client, arguments, handler):
for line in lines: for line in lines:
if ":" in line: if ":" in line:
cmd = line.split(":")[0] cmd = line.split(":")[0]
marguments = line.split(":")[1]
try: try:
opts = parse_args(cmd, marguments, arguments) opts = parse_args(cmd, line.split(":")[1], arguments)
except Exception: except Exception:
response = "Wrong arguments\n<NAK>\n" response = "Wrong arguments\n<NAK>\n"
opts = None opts = None
if opts is not None: if opts is not None:
if handler(cmd, opts): response = "<ACK>\n" if handler(cmd, opts) else "<NAK>\n"
response = "<ACK>\n"
else:
response = "<NAK>\n"
connection.sendall(bytes(response, 'utf-8')) connection.sendall(bytes(response, 'utf-8'))
finally: finally:
connection.close() connection.close()
@ -341,11 +337,7 @@ def do_tcp_server(client, arguments, handler):
def parse_args(cmd, args, mainargs): def parse_args(cmd, args, mainargs):
options = {} options = {}
opts = None opts = args.split(",") if "," in args else [args]
if "," in args:
opts = args.split(",")
else:
opts = [args]
for arg in mainargs: for arg in mainargs:
if "--" in arg: if "--" in arg:
options[arg] = mainargs[arg] options[arg] = mainargs[arg]
@ -605,7 +597,7 @@ class elf:
elif self.elfclass == 2: # 64Bit elif self.elfclass == 2: # 64Bit
start = 0x34 start = 0x34
else: else:
print("Error on parsing " + self.filename) print(f"Error on parsing {self.filename}")
return ['', ''] return ['', '']
elfheadersize, programheaderentrysize, programheaderentrycount = struct.unpack("<HHH", elfheadersize, programheaderentrysize, programheaderentrycount = struct.unpack("<HHH",
self.data[start:start + 3 * 2]) self.data[start:start + 3 * 2])
@ -821,7 +813,7 @@ def print_progress(iteration, total, prefix='', suffix='', decimals=1, bar_lengt
filled_length = int(round(bar_length * iteration / float(total))) filled_length = int(round(bar_length * iteration / float(total)))
bar = '' * filled_length + '-' * (bar_length - filled_length) bar = '' * filled_length + '-' * (bar_length - filled_length)
sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percents, '%', suffix)) sys.stdout.write(f'\r{prefix} |{bar}| {percents}% {suffix}')
if iteration == total: if iteration == total:
sys.stdout.write('\n') sys.stdout.write('\n')