IO Operations

IOContext

A context class for managing IO connections (local, remote, SSH, Docker, kernel).

Attributes:

Name	Type	Description
args		CLI arguments for the connection.
config		Configuration for the target (e.g., binary path, environment).
ssh_conn	Any | None	The SSH connection object.
conn	Any | None	The active connection object (e.g., process, remote, ssh).
proc	Any | None	The process object for local debugging.

Source code in src/pwninit/io.py

class IOContext:
    """
    A context class for managing IO connections (local, remote, SSH, Docker, kernel).

    Attributes:
        args: CLI arguments for the connection.
        config: Configuration for the target (e.g., binary path, environment).
        ssh_conn: The SSH connection object.
        conn: The active connection object (e.g., `process`, `remote`, `ssh`).
        proc: The process object for local debugging.
    """

    def __init__(
        self,
        args: Any,
        config: Any,
        proc: Any | None = None,
        conn: Any | None = None,
        ssh_conn: Any | None = None,
    ) -> None:
        self.args = args
        self.config = config
        self.ssh_conn: Any | None = ssh_conn
        self.conn: Any | None = conn
        self.proc: Any | None = proc

    def __getattr__(self, name: str) -> Any:
        if name != "conn" and self.conn is not None:
            return getattr(self.conn, name)
        raise AttributeError(name)

    def __create_remote_connection(self) -> Any:
        return remote(
            self.args.remote.host,
            self.args.remote.port,
            ssl=self.args.ssl,
        )

    def __create_ssh_connection(self) -> Any:
        return ssh(
            user=self.args.remote.user,
            password=self.args.remote.password,
            host=self.args.remote.host,
            port=self.args.remote.port,
        )

    def __create_ssh_process(self) -> Any:
        if self.args.debug:
            return gdb.debug(self.config.chall, ssh=self.ssh_conn, cwd=self.args.remote.path)
        else:
            return self.ssh_conn.process(self.config.chall, env=self.config.env, cwd=self.args.remote.path)

    def __create_kernel_process(self) -> Any:
        status = log.progress("compiling exploit")
        subprocess.run(
            ["make"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        status.success("done")

        status = log.progress("injecting exploit")
        if not inject(self.config.archive, "exploit"):
            status.failure("failed")
        status.success("done")

        gdb_script = ""
        if self.args.debug:
            gdb_script = self.args.gdb_cmd if self.args.gdb_cmd else ""
            self.config.chall.append("-s")

        p = process(self.config.chall, env=self.config.env)

        if self.args.debug:
            gdb.attach(
                target=("localhost", 1234),
                exe=f"{self.config.kernel}.elf",
                gdbscript=gdb_script,
            )

        return p

    def __create_local_process(self) -> Any:
        if self.config.archive:
            return self.__create_kernel_process()

        gdb_script = self.args.gdb_cmd if self.args.gdb_cmd else ""
        if self.args.debug:
            return gdb.debug(
                self.config.chall,
                gdbscript=gdb_script,
                env=self.config.env,
            )
        elif self.args.strace:
            return process(
                ["strace", "-o", "strace.out", self.config.chall],
                env=self.config.env,
            )
        else:
            p = process(self.config.chall, env=self.config.env)
            if self.args.attach:
                gdb.attach(p, gdbscript=gdb_script)
                log.info("Attached gdb")
                pause()
            return p

    def __launch_docker(self) -> Any:
        client = docker.from_env()
        name = Path(".").resolve().name
        image_tag = getattr(
            self.config, "docker_image", f"pwninit-{name}:latest"
        ).lower()

        container = next((c for c in client.containers.list() if c.image.tags[-1] == image_tag), None)

        if not container:
            try:
                container = client.containers.run(
                    image_tag,
                    pid_mode="host",
                    ports={
                        f"{self.args.remote.port}/tcp": self.args.remote.port
                    },
                    privileged=True,
                    detach=True,
                )
            except docker.errors.APIError as e:
                log.warning(f"Failed to launch docker: {e}")
                exit(1)

        return container

    def __debug_docker(self, container: Any) -> None:
        processes = container.top()
        pid: int | None = None
        bin: str | None = None

        if hasattr(self.config, "docker_bin"):
            bin = self.config.docker_bin
        else:
            for p in processes["Processes"]:
                if self.config.binary in p[-1]:
                    pid = int(p[1])
                    break
                if "socat" in p[-1]:
                    bin = p[-1].split("exec:")[1].split(",")[0]
                    break

            if not bin and not pid:
                log.warning(
                    "No socat running nor binary — set docker_bin in Config"
                )
                exit(1)

        if not pid:
            for p in processes["Processes"]:
                if bin in p[-1]:
                    pid = int(p[1])

        if not pid:
            log.warning("Bin isn't running — check Dockerfile or bin name")
            exit(1)

        gdb.attach(pid, exe=self.config.binary)

    def test_connection(self) -> bool:
        try:
            buf = self.recv(timeout=2)
            self.unrecv(buf)
            return True
        except EOFError:
            log.warning("Failed to connect to docker")
            return False

    def connect(self, enable_log: bool = True) -> "IOContext | None":
        """
        Establish a connection based on the provided arguments and configuration.

        Args:
            enable_log (bool): If True, enable logging (default: True).

        Returns:
            self: On success.
            None: On failure.
        """
        if not enable_log:
            log_level = context.log_level
            context.log_level = "error"

        if self.conn:
            return self

        is_local_process = not self.args.remote or (
            self.args.local and not self.proc
        )
        is_docker_debug = self.args.docker and (
            self.args.debug or self.args.attach
        )
        is_ssh = isinstance(self.args.remote, SSH)

        if is_local_process:
            self.conn = self.proc = self.__create_local_process()

        if self.args.docker:
            container = self.__launch_docker()

        if self.args.remote and is_ssh:
            self.ssh_conn = self.__create_ssh_connection()
            if not self.ssh_conn:
                return None
            self.conn = self.__create_ssh_process()
        elif self.args.remote:
            self.conn = self.__create_remote_connection()

        if is_docker_debug and self.test_connection():
            self.__debug_docker(container)

        if not self.conn:
            log.warning("Failed to create process")
            return None

        if not enable_log:
            context.log_level = log_level

        return self

    def reconnect(self, enable_log: bool = True) -> "IOContext | None":
        """
        Close the current connection and reconnect.

        Args:
            enable_log (bool): If True, enable logging (default: True).

        Returns:
            IOContext | None: The reconnected context, or None on failure.
        """
        if self.conn:
            self.close(enable_log)
        return self.connect(enable_log)

    def close(self, enable_log: bool = True) -> None:
        """
        Close the current connection.

        Args:
            enable_log (bool): If True, enable logging (default: True).
        """
        if not self.conn:
            return
        if not enable_log:
            log_level = context.log_level
            context.log_level = "error"
        self.conn.close()
        self.conn = None
        if not enable_log:
            context.log_level = log_level

    def prompt(self, data: str | bytes, **kwargs: Any) -> None:
        """
        Send data to the target, optionally waiting for a prefix.

        Args:
            data: The data to send.
            **kwargs: Additional arguments for `send`/`sendline` (e.g., `prefix`, `line`).
        """
        prefix = utils.encode(kwargs.pop("prefix", self.config.prefix))
        line = utils.encode(kwargs.pop("line", True))
        data = utils.encode(data)
        r = self.conn

        if prefix and line:
            r.sendlineafter(prefix, data, **kwargs)
        elif prefix:
            r.sendafter(prefix, data, **kwargs)
        elif not prefix and line:
            r.sendline(data, **kwargs)
        else:
            r.send(data, **kwargs)

    def sla(self, *args: str | bytes, **kwargs: Any) -> None:
        """
        Send a line after a prefix (shorthand for `prompt` with `line=True`).

        Args:
            *args: If one argument, it is the data to send. If two, the first is the prefix and the second is the data.
            **kwargs: Additional arguments for `prompt`.
        """
        if len(args) == 1:
            self.prompt(args[0], **kwargs)
        elif len(args) >= 2:
            self.prompt(args[1], prefix=args[0], **kwargs)

    def sa(self, *args: str | bytes, **kwargs: Any) -> None:
        """
        Send data after a prefix (shorthand for `prompt` with `line=False`).

        Args:
            *args: If one argument, it is the data to send. If two, the first is the prefix and the second is the data.
            **kwargs: Additional arguments for `prompt`.
        """
        self.sla(*args, line=False, **kwargs)

    def sl(self, data: str | bytes, **kwargs: Any) -> None:
        """
        Send a line without waiting for a prefix.

        Args:
            data: The data to send.
            **kwargs: Additional arguments for `prompt`.
        """
        self.prompt(data, prefix=None, **kwargs)

    def send(self, data: str | bytes, **kwargs: Any) -> None:
        """
        Send data without waiting for a prefix or newline.

        Args:
            data: The data to send.
            **kwargs: Additional arguments for `prompt`.
        """
        self.prompt(data, prefix=None, line=False, **kwargs)

    def recv(
        self,
        prefix: str | bytes | int | None = None,
        **kwargs: Any,
    ) -> bytes:
        """
        Receive data from the target, optionally waiting for a prefix.

        Args:
            prefix: The prefix to wait for (bytes, str, or int for `recvn`).
            **kwargs: Additional arguments for `recv`/`recvline`/`recvuntil`.

        Returns:
            bytes: The received data.
        """
        r = kwargs.pop("io", self.conn)
        line = kwargs.pop("line", False)

        if prefix is None:
            return r.recvline(**kwargs) if line else r.recv(**kwargs)
        elif isinstance(prefix, int):
            return r.recvn(prefix, **kwargs)

        prefix = utils.encode(prefix)
        drop = kwargs.pop("drop", True)
        if line:
            r.recvuntil(prefix, drop=drop, **kwargs)
            return r.recvline(drop=drop, **kwargs)
        else:
            return r.recvuntil(prefix, drop=drop, **kwargs)

    def ru(self, u: str | bytes, **kwargs: Any) -> bytes:
        """
        Receive data until a specific string (shorthand for `recv` with `prefix`).

        Args:
            u: The string to wait for.
            **kwargs: Additional arguments for `recv`.
        """
        return self.recv(u, **kwargs)

    def rl(self, **kwargs: Any) -> bytes:
        """
        Receive a line (shorthand for `recv` with `line=True`).

        Returns:
            bytes: The received line.
        """
        return self.recv(line=True, **kwargs)

    def rla(self, d: str | bytes, **kwargs: Any) -> bytes:
        """
        Receive a line after a specific string (shorthand for `recv` with `prefix` and `line=True`).

        Args:
            d: The string to wait for.
            **kwargs: Additional arguments for `recv`.
        """
        return self.recv(d, line=True, **kwargs)

    def ra(self) -> bytes:
        """
        Receive all remaining data.

        Returns:
            bytes: All remaining data.
        """
        return self.recvall()

    def itrv(self) -> Any:
        """
        Switch to interactive mode.

        Returns:
            Any: The result of the interactive session.
        """
        return self.interactive()

    def urecv(self) -> bytes:
        """
        Undo the last receive operation.

        Returns:
            bytes: The data that was un-received.
        """
        return self.unrecv()

    def rln(self, n: int, **kwargs: Any) -> list[bytes]:
        """
        Receive `n` lines.

        Args:
            n (int): The number of lines to receive.
            **kwargs: Additional arguments for `recv`.

        Returns:
            list[bytes]: A list of received lines.
        """
        return [self.rl(**kwargs) for _ in range(n)]

    def pow(self) -> None:
        """
        Solve any pows implemented in utils using what's been received.
        """
        utils.solve_pow(self.clean())

close(enable_log=True)

Close the current connection.

Parameters:

Name	Type	Description	Default
enable_log	bool	If True, enable logging (default: True).	True

Source code in src/pwninit/io.py

def close(self, enable_log: bool = True) -> None:
    """
    Close the current connection.

    Args:
        enable_log (bool): If True, enable logging (default: True).
    """
    if not self.conn:
        return
    if not enable_log:
        log_level = context.log_level
        context.log_level = "error"
    self.conn.close()
    self.conn = None
    if not enable_log:
        context.log_level = log_level

connect(enable_log=True)

Establish a connection based on the provided arguments and configuration.

Parameters:

Name	Type	Description	Default
enable_log	bool	If True, enable logging (default: True).	True

Returns:

Name	Type	Description
self	IOContext | None	On success.
None	IOContext | None	On failure.

Source code in src/pwninit/io.py

def connect(self, enable_log: bool = True) -> "IOContext | None":
    """
    Establish a connection based on the provided arguments and configuration.

    Args:
        enable_log (bool): If True, enable logging (default: True).

    Returns:
        self: On success.
        None: On failure.
    """
    if not enable_log:
        log_level = context.log_level
        context.log_level = "error"

    if self.conn:
        return self

    is_local_process = not self.args.remote or (
        self.args.local and not self.proc
    )
    is_docker_debug = self.args.docker and (
        self.args.debug or self.args.attach
    )
    is_ssh = isinstance(self.args.remote, SSH)

    if is_local_process:
        self.conn = self.proc = self.__create_local_process()

    if self.args.docker:
        container = self.__launch_docker()

    if self.args.remote and is_ssh:
        self.ssh_conn = self.__create_ssh_connection()
        if not self.ssh_conn:
            return None
        self.conn = self.__create_ssh_process()
    elif self.args.remote:
        self.conn = self.__create_remote_connection()

    if is_docker_debug and self.test_connection():
        self.__debug_docker(container)

    if not self.conn:
        log.warning("Failed to create process")
        return None

    if not enable_log:
        context.log_level = log_level

    return self

itrv()

Switch to interactive mode.

Returns:

Name	Type	Description
Any	Any	The result of the interactive session.

Source code in src/pwninit/io.py

def itrv(self) -> Any:
    """
    Switch to interactive mode.

    Returns:
        Any: The result of the interactive session.
    """
    return self.interactive()

pow()

Solve any pows implemented in utils using what's been received.

Source code in src/pwninit/io.py

def pow(self) -> None:
    """
    Solve any pows implemented in utils using what's been received.
    """
    utils.solve_pow(self.clean())

prompt(data, **kwargs)

Send data to the target, optionally waiting for a prefix.

Parameters:

Name	Type	Description	Default
data	str | bytes	The data to send.	required
**kwargs	Any	Additional arguments for send/sendline (e.g., prefix, line).	{}

Source code in src/pwninit/io.py

def prompt(self, data: str | bytes, **kwargs: Any) -> None:
    """
    Send data to the target, optionally waiting for a prefix.

    Args:
        data: The data to send.
        **kwargs: Additional arguments for `send`/`sendline` (e.g., `prefix`, `line`).
    """
    prefix = utils.encode(kwargs.pop("prefix", self.config.prefix))
    line = utils.encode(kwargs.pop("line", True))
    data = utils.encode(data)
    r = self.conn

    if prefix and line:
        r.sendlineafter(prefix, data, **kwargs)
    elif prefix:
        r.sendafter(prefix, data, **kwargs)
    elif not prefix and line:
        r.sendline(data, **kwargs)
    else:
        r.send(data, **kwargs)

ra()

Receive all remaining data.

Returns:

Name	Type	Description
bytes	bytes	All remaining data.

Source code in src/pwninit/io.py

def ra(self) -> bytes:
    """
    Receive all remaining data.

    Returns:
        bytes: All remaining data.
    """
    return self.recvall()

reconnect(enable_log=True)

Close the current connection and reconnect.

Parameters:

Name	Type	Description	Default
enable_log	bool	If True, enable logging (default: True).	True

Returns:

Type	Description
IOContext | None	IOContext | None: The reconnected context, or None on failure.

Source code in src/pwninit/io.py

def reconnect(self, enable_log: bool = True) -> "IOContext | None":
    """
    Close the current connection and reconnect.

    Args:
        enable_log (bool): If True, enable logging (default: True).

    Returns:
        IOContext | None: The reconnected context, or None on failure.
    """
    if self.conn:
        self.close(enable_log)
    return self.connect(enable_log)

recv(prefix=None, **kwargs)

Receive data from the target, optionally waiting for a prefix.

Parameters:

Name	Type	Description	Default
prefix	str | bytes | int | None	The prefix to wait for (bytes, str, or int for recvn).	None
**kwargs	Any	Additional arguments for recv/recvline/recvuntil.	{}

Returns:

Name	Type	Description
bytes	bytes	The received data.

Source code in src/pwninit/io.py

def recv(
    self,
    prefix: str | bytes | int | None = None,
    **kwargs: Any,
) -> bytes:
    """
    Receive data from the target, optionally waiting for a prefix.

    Args:
        prefix: The prefix to wait for (bytes, str, or int for `recvn`).
        **kwargs: Additional arguments for `recv`/`recvline`/`recvuntil`.

    Returns:
        bytes: The received data.
    """
    r = kwargs.pop("io", self.conn)
    line = kwargs.pop("line", False)

    if prefix is None:
        return r.recvline(**kwargs) if line else r.recv(**kwargs)
    elif isinstance(prefix, int):
        return r.recvn(prefix, **kwargs)

    prefix = utils.encode(prefix)
    drop = kwargs.pop("drop", True)
    if line:
        r.recvuntil(prefix, drop=drop, **kwargs)
        return r.recvline(drop=drop, **kwargs)
    else:
        return r.recvuntil(prefix, drop=drop, **kwargs)

rl(**kwargs)

Receive a line (shorthand for recv with line=True).

Returns:

Name	Type	Description
bytes	bytes	The received line.

Source code in src/pwninit/io.py

def rl(self, **kwargs: Any) -> bytes:
    """
    Receive a line (shorthand for `recv` with `line=True`).

    Returns:
        bytes: The received line.
    """
    return self.recv(line=True, **kwargs)

rla(d, **kwargs)

Receive a line after a specific string (shorthand for recv with prefix and line=True).

Parameters:

Name	Type	Description	Default
d	str | bytes	The string to wait for.	required
**kwargs	Any	Additional arguments for recv.	{}

Source code in src/pwninit/io.py

def rla(self, d: str | bytes, **kwargs: Any) -> bytes:
    """
    Receive a line after a specific string (shorthand for `recv` with `prefix` and `line=True`).

    Args:
        d: The string to wait for.
        **kwargs: Additional arguments for `recv`.
    """
    return self.recv(d, line=True, **kwargs)

rln(n, **kwargs)

Receive n lines.

Parameters:

Name	Type	Description	Default
n	int	The number of lines to receive.	required
**kwargs	Any	Additional arguments for recv.	{}

Returns:

Type	Description
list[bytes]	list[bytes]: A list of received lines.

Source code in src/pwninit/io.py

def rln(self, n: int, **kwargs: Any) -> list[bytes]:
    """
    Receive `n` lines.

    Args:
        n (int): The number of lines to receive.
        **kwargs: Additional arguments for `recv`.

    Returns:
        list[bytes]: A list of received lines.
    """
    return [self.rl(**kwargs) for _ in range(n)]

ru(u, **kwargs)

Receive data until a specific string (shorthand for recv with prefix).

Parameters:

Name	Type	Description	Default
u	str | bytes	The string to wait for.	required
**kwargs	Any	Additional arguments for recv.	{}

Source code in src/pwninit/io.py

def ru(self, u: str | bytes, **kwargs: Any) -> bytes:
    """
    Receive data until a specific string (shorthand for `recv` with `prefix`).

    Args:
        u: The string to wait for.
        **kwargs: Additional arguments for `recv`.
    """
    return self.recv(u, **kwargs)

sa(*args, **kwargs)

Send data after a prefix (shorthand for prompt with line=False).

Parameters:

Name	Type	Description	Default
*args	str | bytes	If one argument, it is the data to send. If two, the first is the prefix and the second is the data.	()
**kwargs	Any	Additional arguments for prompt.	{}

Source code in src/pwninit/io.py

def sa(self, *args: str | bytes, **kwargs: Any) -> None:
    """
    Send data after a prefix (shorthand for `prompt` with `line=False`).

    Args:
        *args: If one argument, it is the data to send. If two, the first is the prefix and the second is the data.
        **kwargs: Additional arguments for `prompt`.
    """
    self.sla(*args, line=False, **kwargs)

send(data, **kwargs)

Send data without waiting for a prefix or newline.

Parameters:

Name	Type	Description	Default
data	str | bytes	The data to send.	required
**kwargs	Any	Additional arguments for prompt.	{}

Source code in src/pwninit/io.py

def send(self, data: str | bytes, **kwargs: Any) -> None:
    """
    Send data without waiting for a prefix or newline.

    Args:
        data: The data to send.
        **kwargs: Additional arguments for `prompt`.
    """
    self.prompt(data, prefix=None, line=False, **kwargs)

sl(data, **kwargs)

Send a line without waiting for a prefix.

Parameters:

Name	Type	Description	Default
data	str | bytes	The data to send.	required
**kwargs	Any	Additional arguments for prompt.	{}

Source code in src/pwninit/io.py

def sl(self, data: str | bytes, **kwargs: Any) -> None:
    """
    Send a line without waiting for a prefix.

    Args:
        data: The data to send.
        **kwargs: Additional arguments for `prompt`.
    """
    self.prompt(data, prefix=None, **kwargs)

sla(*args, **kwargs)

Send a line after a prefix (shorthand for prompt with line=True).

Parameters:

Name	Type	Description	Default
*args	str | bytes	If one argument, it is the data to send. If two, the first is the prefix and the second is the data.	()
**kwargs	Any	Additional arguments for prompt.	{}

Source code in src/pwninit/io.py

def sla(self, *args: str | bytes, **kwargs: Any) -> None:
    """
    Send a line after a prefix (shorthand for `prompt` with `line=True`).

    Args:
        *args: If one argument, it is the data to send. If two, the first is the prefix and the second is the data.
        **kwargs: Additional arguments for `prompt`.
    """
    if len(args) == 1:
        self.prompt(args[0], **kwargs)
    elif len(args) >= 2:
        self.prompt(args[1], prefix=args[0], **kwargs)

urecv()

Undo the last receive operation.

Returns:

Name	Type	Description
bytes	bytes	The data that was un-received.

Source code in src/pwninit/io.py

def urecv(self) -> bytes:
    """
    Undo the last receive operation.

    Returns:
        bytes: The data that was un-received.
    """
    return self.unrecv()

NC dataclass

Dataclass for representing a network connection.

Attributes:

Name	Type	Description
host	str	The host address.
port	int	The port number.

Source code in src/pwninit/io.py

@dataclass(frozen=True)
class NC:
    """
    Dataclass for representing a network connection.

    Attributes:
        host (str): The host address.
        port (int): The port number.
    """

    host: str
    port: int

SSH dataclass

Dataclass for representing an SSH connection.

Attributes:

Name	Type	Description
user	str	The SSH username.
host	str	The host address.
password	str	The SSH password (default: "").
port	int	The SSH port (default: 22).

Source code in src/pwninit/io.py

@dataclass(frozen=True)
class SSH:
    """
    Dataclass for representing an SSH connection.

    Attributes:
        user (str): The SSH username.
        host (str): The host address.
        password (str): The SSH password (default: "").
        port (int): The SSH port (default: 22).
    """

    user: str
    host: str
    password: str = ""
    port: int = 22
    path: str = "."