rz-pipe/python/rzpipe/open_async.py

# -*- coding: utf-8 -*-
"""open_async.py
This script use code from rzpipe-async/open_p3.py script.

"""
import asyncio
import os
import re


# whole file doesn't have any profit of asyncio usage, TODO: refactor

from collections.abc import Iterable
from contextlib import ContextDecorator
from urllib.parse import quote, urlparse

from .open_base import OpenBase, get_rizin_path


class open(OpenBase, ContextDecorator):
    # --------------------------------------------------------------------------
    # Contenxt manager functions
    # --------------------------------------------------------------------------
    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()
        return False

    def close(self):
        if self._loop.is_running():
            self._loop.stop()
        if not self._loop.is_closed():
            self._loop.close()

    def __init__(self, filename="", flags=[], rizinhome=None):
        super(open, self).__init__(filename, flags)

        self.rzhome = rizinhome

        if os.name == "nt":
            self._loop = asyncio.ProactorEventLoop()
            asyncio.set_event_loop(self._loop)
        else:
            watcher = asyncio.get_child_watcher()
            self._loop = asyncio.new_event_loop()
            watcher.attach_loop(self._loop)

        self.asyn = True

        if filename.startswith("http://"):
            self._cmd_coro = self._cmd_http
            self.uri = "/cmd"

            _tmp = urlparse(filename)
            self._host = _tmp.hostname
            self._port = _tmp.port

        elif filename.startswith("ccall://"):
            self._cmd_coro = self._cmd_native
            self.uri = filename[7:]

        elif filename.startswith("tcp://"):

            r = re.match(r"tcp://(\d+\.\d+.\d+.\d+):(\d+)/?", filename)
            if not r:
                raise Exception("String doesn't match tcp format")

            self._cmd_coro = self._cmd_tcp
            self._host = r.group(1)
            self._port = r.group(2)

        elif filename:

            self._cmd_coro = self._cmd_process

            cmd = ["-q0", filename]
            cmd = cmd[:1] + flags + cmd[1:]
            self._process_start_cmd = cmd

        else:
            self.asyn = False

    def _callback_wrapper(self, future):
        result, callback = future.result()

        if callback:
            callback(result)

    def _cmd(self, cmd, **kwargs):
        # Get callback, if available
        callback = kwargs.get("callback")
        future = asyncio.Future(loop=self._loop)
        future.add_done_callback(self._callback_wrapper)

        task = self._loop.create_task(self._cmd_coro(cmd, future, callback))

        # Create and start a new task (coroutine)
        self._loop.run_until_complete(task)
        return task.result() if task else None

    @asyncio.coroutine
    def _cmd_process(self, cmd, future, callback):
        if not hasattr(self, "process"):
            if self.rzhome is not None:
                if not os.path.isdir(self.rzhome):
                    raise Exception(
                        "`rizinhome` passed to `open` is invalid, leave it None or put a valid path to rizin folder"
                    )
                rzpath = os.path.join(self.rzhome, "rizin")
                if os.name == "nt":
                    rzpath += ".exe"
            else:
                rzpath = get_rizin_path()

            create = asyncio.create_subprocess_exec(
                rzpath,
                *self._process_start_cmd,
                shell=False,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                loop=self._loop
            )

            self.process = yield from create  # Init the process

            yield from self.process.stdout.read(1)  # Reads initial \x00

        cmd = cmd.strip().replace("\n", ";")
        self.process.stdin.write(bytes(cmd + "\n", "utf-8"))

        out = []
        while True:
            # foo = self.process.stdout.read(1)
            foo = yield from self.process.stdout.read(1)
            if foo == b"\x00":
                break
            if len(foo) < 1:
                return None
            out.append(foo)

        out = b"".join(out).decode("utf-8")
        future.set_result((out, callback))
        return out

    @asyncio.coroutine
    def _cmd_http(self, cmd, future, callback):
        try:
            quocmd = quote(cmd)

            reader, writer = yield from asyncio.open_connection(
                self._host, self._port, loop=self._loop
            )

            message = "\n\r".join(
                [
                    "GET /cmd/%s HTTP/1.1" % quocmd,
                    "Host: %s:%s" % (self._host, self._port),
                    "User-Agent: rzpipe/Python Client",
                    "Accept: */*",
                    "",
                    "",
                ]
            ).encode()

            writer.write(message)
            data = yield from reader.read(512)
            res = [data]
            while data:
                data = yield from reader.read(512)
                res.append(data)
            writer.close()

            res = b"".join(res)

            # Remove http headers
            start = 0
            for x in res.splitlines():
                if not x:
                    start += 1
                    break
                start += len(x) + 1  # +1 because we must be count '\n'
            res = res[start:].decode("utf-8")
            future.set_result((res, callback))
            return res

        except Exception as e:
            future.set_result((str(e), callback))

    @asyncio.coroutine
    def _cmd_tcp(self, cmd, future, callback):

        try:
            reader, writer = yield from asyncio.open_connection(
                self._host, self._port, loop=self._loop
            )

            writer.write(cmd.encode("utf-8"))
            data = yield from reader.read(512)

            res = [data]
            while data:
                res.append(data)
                data = yield from reader.read(512)
            res = b"".join(res).decode("utf-8")
            future.set_result((res, callback))
            writer.close()
            return res

        except Exception as e:
            future.set_result((str(e), callback))

    def wait(self, task):
        """Wait until task finish"""
        _tasks = task
        if not isinstance(task, Iterable):
            _tasks = [task]

        if self._loop.is_running():
            asyncio.wait(_tasks, loop=self._loop)