209 lines
7.0 KiB
Python
209 lines
7.0 KiB
Python
import argparse
|
|
import codecs
|
|
import os
|
|
import sys
|
|
import time
|
|
from threading import Event, Thread
|
|
from typing import AnyStr, List, MutableMapping, Optional
|
|
|
|
import frida
|
|
from colorama import Fore, Style
|
|
|
|
from frida_tools.application import ConsoleApplication
|
|
from frida_tools.stream_controller import DisposedException, StreamController
|
|
from frida_tools.units import bytes_to_megabytes
|
|
|
|
|
|
def main() -> None:
|
|
app = PushApplication()
|
|
app.run()
|
|
|
|
|
|
class PushApplication(ConsoleApplication):
|
|
def _add_options(self, parser: argparse.ArgumentParser) -> None:
|
|
parser.add_argument("files", help="local files to push", nargs="+")
|
|
|
|
def _usage(self) -> str:
|
|
return "%(prog)s [options] LOCAL... REMOTE"
|
|
|
|
def _initialize(self, parser: argparse.ArgumentParser, options: argparse.Namespace, args: List[str]) -> None:
|
|
paths = options.files
|
|
if len(paths) == 1:
|
|
raise ValueError("missing remote path")
|
|
self._local_paths = paths[:-1]
|
|
self._remote_path = paths[-1]
|
|
|
|
self._script: Optional[frida.core.Script] = None
|
|
self._stream_controller: Optional[StreamController] = None
|
|
self._total_bytes = 0
|
|
self._time_started: Optional[float] = None
|
|
self._completed = Event()
|
|
self._transfers: MutableMapping[str, bool] = {}
|
|
|
|
def _needs_target(self) -> bool:
|
|
return False
|
|
|
|
def _start(self) -> None:
|
|
try:
|
|
self._attach(0)
|
|
|
|
data_dir = os.path.dirname(__file__)
|
|
with codecs.open(os.path.join(data_dir, "fs_agent.js"), "r", "utf-8") as f:
|
|
source = f.read()
|
|
|
|
def on_message(message, data) -> None:
|
|
self._reactor.schedule(lambda: self._on_message(message, data))
|
|
|
|
assert self._session is not None
|
|
script = self._session.create_script(name="push", source=source)
|
|
self._script = script
|
|
script.on("message", on_message)
|
|
self._on_script_created(script)
|
|
script.load()
|
|
|
|
self._stream_controller = StreamController(
|
|
self._post_stream_stanza, on_stats_updated=self._on_stream_stats_updated
|
|
)
|
|
|
|
worker = Thread(target=self._perform_push)
|
|
worker.start()
|
|
except Exception as e:
|
|
self._update_status(f"Failed to push: {e}")
|
|
self._exit(1)
|
|
return
|
|
|
|
def _stop(self) -> None:
|
|
for path in self._local_paths:
|
|
if path not in self._transfers:
|
|
self._complete_transfer(path, success=False)
|
|
|
|
if self._stream_controller is not None:
|
|
self._stream_controller.dispose()
|
|
|
|
def _perform_push(self) -> None:
|
|
for path in self._local_paths:
|
|
try:
|
|
self._total_bytes += os.path.getsize(path)
|
|
except:
|
|
pass
|
|
self._time_started = time.time()
|
|
|
|
for i, path in enumerate(self._local_paths):
|
|
filename = os.path.basename(path)
|
|
|
|
try:
|
|
with open(path, "rb") as f:
|
|
assert self._stream_controller is not None
|
|
sink = self._stream_controller.open(str(i), {"filename": filename, "target": self._remote_path})
|
|
while True:
|
|
chunk = f.read(4 * 1024 * 1024)
|
|
if len(chunk) == 0:
|
|
break
|
|
sink.write(chunk)
|
|
sink.close()
|
|
except DisposedException:
|
|
break
|
|
except Exception as e:
|
|
self._print_error(str(e))
|
|
self._complete_transfer(path, success=False)
|
|
|
|
self._completed.wait()
|
|
|
|
self._reactor.schedule(lambda: self._on_push_finished())
|
|
|
|
def _on_push_finished(self) -> None:
|
|
successes = self._transfers.values()
|
|
|
|
if any(successes):
|
|
self._render_summary_ui()
|
|
|
|
status = 0 if all(successes) else 1
|
|
self._exit(status)
|
|
|
|
def _render_progress_ui(self) -> None:
|
|
if self._completed.is_set():
|
|
return
|
|
assert self._stream_controller is not None
|
|
megabytes_sent = bytes_to_megabytes(self._stream_controller.bytes_sent)
|
|
total_megabytes = bytes_to_megabytes(self._total_bytes)
|
|
if total_megabytes != 0 and megabytes_sent <= total_megabytes:
|
|
self._update_status(f"Pushed {megabytes_sent:.1f} out of {total_megabytes:.1f} MB")
|
|
else:
|
|
self._update_status(f"Pushed {megabytes_sent:.1f} MB")
|
|
|
|
def _render_summary_ui(self) -> None:
|
|
assert self._time_started is not None
|
|
duration = time.time() - self._time_started
|
|
|
|
if len(self._local_paths) == 1:
|
|
prefix = f"{self._local_paths[0]}: "
|
|
else:
|
|
prefix = ""
|
|
|
|
files_transferred = sum(map(int, self._transfers.values()))
|
|
|
|
assert self._stream_controller is not None
|
|
bytes_sent = self._stream_controller.bytes_sent
|
|
megabytes_per_second = bytes_to_megabytes(bytes_sent) / duration
|
|
|
|
self._update_status(
|
|
"{}{} file{} pushed. {:.1f} MB/s ({} bytes in {:.3f}s)".format(
|
|
prefix,
|
|
files_transferred,
|
|
"s" if files_transferred != 1 else "",
|
|
megabytes_per_second,
|
|
bytes_sent,
|
|
duration,
|
|
)
|
|
)
|
|
|
|
def _on_message(self, message, data) -> None:
|
|
handled = False
|
|
|
|
if message["type"] == "send":
|
|
payload = message["payload"]
|
|
ptype = payload["type"]
|
|
if ptype == "stream":
|
|
stanza = payload["payload"]
|
|
self._stream_controller.receive(stanza, data)
|
|
handled = True
|
|
elif ptype == "push:io-success":
|
|
index = payload["index"]
|
|
self._on_io_success(self._local_paths[index])
|
|
handled = True
|
|
elif ptype == "push:io-error":
|
|
index = payload["index"]
|
|
self._on_io_error(self._local_paths[index], payload["error"])
|
|
handled = True
|
|
|
|
if not handled:
|
|
self._print(message)
|
|
|
|
def _on_io_success(self, local_path: str) -> None:
|
|
self._complete_transfer(local_path, success=True)
|
|
|
|
def _on_io_error(self, local_path: str, error) -> None:
|
|
self._print_error(f"{local_path}: {error}")
|
|
self._complete_transfer(local_path, success=False)
|
|
|
|
def _complete_transfer(self, local_path: str, success: bool) -> None:
|
|
self._transfers[local_path] = success
|
|
if len(self._transfers) == len(self._local_paths):
|
|
self._completed.set()
|
|
|
|
def _post_stream_stanza(self, stanza, data: Optional[AnyStr] = None) -> None:
|
|
self._script.post({"type": "stream", "payload": stanza}, data=data)
|
|
|
|
def _on_stream_stats_updated(self) -> None:
|
|
self._render_progress_ui()
|
|
|
|
def _print_error(self, message: str) -> None:
|
|
self._print(Fore.RED + Style.BRIGHT + message + Style.RESET_ALL, file=sys.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
pass
|