def test_simple_reading(ch: channel.Channel) -> None:
ch.write(b"1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ")
out = ch.read(10)
assert out == b"1234567890"
out = ch.read()
assert out == b"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
def test_taking(ch: channel.Channel) -> None:
ch.sendline("echo Hello")
ch2 = ch.take()
ch2.sendline("echo World")
with pytest.raises(channel.ChannelTakenException):
ch.sendline("echo Illegal")
def test_simple_read_iter(ch: channel.Channel) -> None:
ch.write(b"12345678901234567890")
final = bytearray()
for new in ch.read_iter(10):
final.extend(new)
assert final == b"1234567890"
for i in range(1, 10):
c = ch.read(1)
assert c == str(i).encode("utf-8")
def wait_for_shell(ch: channel.Channel) -> None:
# Repeatedly sends `echo TBOT''LOGIN\r`. At some point, the shell
# interprets this command and prints out `TBOTLOGIN` because of the
# quotation-marks being removed. Once we detect this, this function
# can return, knowing the shell is now running on the other end.
#
# Credit to Pavel for this idea!
while True:
ch.sendline("echo TBOT''LOGIN")
try:
ch.expect("TBOTLOGIN", timeout=0.2)
break
except TimeoutError:
pass
def _ctx(orig_chan: channel.Channel,
cmd_context: CMD_CONTEXT) -> "typing.Iterator[RunCommandProxy]":
"""
Helper function for LinuxShell.run() implementations. See the comment
near the CMD_CONTEXT definition in this file for more details.
"""
with orig_chan.borrow() as ch:
proxy = RunCommandProxy(ch, cmd_context)
try:
yield proxy
except Exception as e:
proxy._cmd_context.throw(e.__class__, e)
proxy._assert_end()
def test_borrowing(ch: channel.Channel) -> None:
ch.sendline("echo Hello")
with ch.borrow() as ch2:
ch2.sendline("echo World")
raised = False
try:
ch.sendline("echo Illegal")
except channel.ChannelBorrowedException:
raised = True
assert raised, "Borrow was unsuccessful"
ch.sendline("echo back again")
def selftest_machine_channel(ch: channel.Channel, remote_close: bool) -> None:
out = ch.raw_command("echo Hello World", timeout=1)
assert out == "Hello World\n", repr(out)
# Check recv_n
ch.send("echo Foo Bar\n")
out2 = ch.recv_n(8, timeout=1.0)
assert out2 == b"echo Foo", repr(out)
ch.read_until_prompt(channel.TBOT_PROMPT)
assert ch.isopen()
if remote_close:
ch.send("exit\n")
time.sleep(0.1)
ch.recv(timeout=1)
raised = False
try:
ch.recv(timeout=1)
except channel.ChannelClosedException:
raised = True
assert raised
else:
ch.close()
assert not ch.isopen()
raised = False
try:
ch.send("\n")
except channel.ChannelClosedException:
raised = True
assert raised
raised = False
try:
ch.recv(timeout=1)
except channel.ChannelClosedException:
raised = True
assert raised
def test_simple_expect3(ch: channel.Channel) -> None:
ch.sendline("echo Lo1337rem")
res = ch.expect(["Dolor", "roloD", tbot.Re(r"Lo(\d{1,20})"), "rem"])
assert res.i == 2
assert isinstance(res.match, Match), "Not a match object"
assert res.match.group(1) == b"1337"
def test_simple_expect2(ch: channel.Channel) -> None:
ch.sendline("echo Lorem Ipsum Dolor Sit")
res = ch.expect(["Lol", "Dolor", "Dol"])
assert res.i == 1
assert res.match == "Dolor"
def test_simple_expect(ch: channel.Channel) -> None:
ch.sendline("echo Lorem Ipsum")
res = ch.expect(["Lol", "Ip"])
assert res.i == 1
assert res.match == "Ip"
def test_simple_readline(ch: channel.Channel) -> None:
ch.sendline("echo Hello; echo World", read_back=True)
out_s = ch.readline()
assert out_s == "Hello\n"
out_s = ch.readline()
assert out_s == "World\n"
def test_simple_command(ch: channel.Channel) -> None:
ch.sendline("echo Hello World", read_back=True)
out = ch.read()
assert out.startswith(b"Hello World")
def __new__(
cls, chan: channel.Channel, cmd_context: CMD_CONTEXT
) -> "RunCommandProxy":
chan.__class__ = cls
return typing.cast(RunCommandProxy, chan)