Skip to content

Commit

Permalink
Ensure readline_expect doesn't leak flags and uses timeout (#70)
Browse files Browse the repository at this point in the history
* ensure readline_expect doesn't leak flags and uses timeout

* fix tests and version bump

Co-authored-by: M:Munier <marc@geomaps.de>
  • Loading branch information
mauricesvp and MMunier authored Jul 13, 2020
1 parent b30e1e8 commit 864d4ae
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 48 deletions.
12 changes: 6 additions & 6 deletions Makefile
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
lint:
python -m isort -c -rc src/ tests/ example/
python -m black --check src/ tests/ example/
python -m flake8 --select F --per-file-ignores="__init__.py:F401" src/ tests/ example/
python -m mypy src/ tests/ example/
python3 -m isort -c -rc src/ tests/ example/
python3 -m black --check src/ tests/ example/
python3 -m flake8 --select F --per-file-ignores="__init__.py:F401" src/ tests/ example/
python3 -m mypy src/ tests/ example/

format:
python -m isort -rc src/ tests/ example/
python -m black src/ tests/ example/
python3 -m isort -rc src/ tests/ example/
python3 -m black src/ tests/ example/

test:
pip install .
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

setuptools.setup(
name="enochecker",
version="0.2.0",
version="0.2.1",
author="domenukk",
author_email="dmaier@sect.tu-berlin.de",
description="Library to build checker scripts for EnoEngine A/D CTF Framework in Python",
Expand Down
1 change: 0 additions & 1 deletion src/enochecker/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
ensure_bytes,
ensure_unicode,
ensure_valid_filename,
readline_expect,
serve_once,
sha256ify,
snake_caseify,
Expand Down
99 changes: 61 additions & 38 deletions src/enochecker/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
Union,
)

from .results import BrokenServiceException, OfflineException
from .results import BrokenServiceException

if TYPE_CHECKING: # pragma: no cover
import requests
Expand Down Expand Up @@ -179,42 +179,42 @@ def debase64ify(
return base64.b64decode(s).decode("utf-8")


def readline_expect(
telnet: Union[telnetlib.Telnet, "SimpleSocket"],
expected: Union[str, bytes],
read_until: Union[str, bytes] = b"\n",
timeout: int = 30,
) -> bytes:
"""
Read to newline (or read_until string) and assert the presence of a string in the response.
Will raise an exception if failed.
:param telnet: Connected telnet instance (the result of self.telnet(..))
:param expected: the expected String to search for in the response
:param read_until: Which char to read until.
:param timeout: a timeout
:return: the bytes read
"""
logger = getattr(telnet, "logger", utilslogger)

if isinstance(expected, str):
expected = expected.encode("utf-8")
if isinstance(read_until, str):
read_until = read_until.encode("utf-8")

read = telnet.read_until(read_until, timeout)
if read == b"":
err = "Expected {!r} but got nothing/timeout!".format(expected)
logger.error(err, stack_info=True)
telnet.close()
raise OfflineException(err)
if expected not in read:
err = "Expected {!r} but got {!r}".format(expected, read)
logger.error(err, stack_info=True)
telnet.close()
raise BrokenServiceException(err)
return read
# def readline_expect(

This comment has been minimized.

Copy link
@domenukk

domenukk Jul 14, 2020

Member

The idea behind this function was that people not using the wrapped socket for whatever reason can still take advantage of this functionality, btw.

# telnet: Union[telnetlib.Telnet, "SimpleSocket"],
# expected: Union[str, bytes],
# read_until: Union[str, bytes] = b"\n",
# timeout: int = 30,
# ) -> bytes:
# """
# Read to newline (or read_until string) and assert the presence of a string in the response.

# Will raise an exception if failed.

# :param telnet: Connected telnet instance (the result of self.telnet(..))
# :param expected: the expected String to search for in the response
# :param read_until: Which char to read until.
# :param timeout: a timeout
# :return: the bytes read
# """
# logger = getattr(telnet, "logger", utilslogger)

# if isinstance(expected, str):
# expected = expected.encode("utf-8")
# if isinstance(read_until, str):
# read_until = read_until.encode("utf-8")

# read = telnet.read_until(read_until, timeout)
# if read == b"":
# err = "Expected {!r} but got nothing/timeout!".format(expected)
# logger.error(err, stack_info=True)
# telnet.close()
# raise OfflineException(err)
# if expected not in read:
# err = "Expected {!r} but got {!r}".format(expected, read)
# logger.error(err, stack_info=True)
# telnet.close()
# raise BrokenServiceException(err)
# return read


def start_daemon(target: Callable[..., Any]) -> threading.Thread:
Expand Down Expand Up @@ -352,6 +352,7 @@ def readline_expect(
expected: Union[str, bytes],
read_until: Union[str, bytes] = b"\n",
timeout: Optional[int] = None,
exception_message: Optional[str] = None,
) -> bytes:
"""
Read to newline (or read_until string) and assert the presence of a string in the response.
Expand All @@ -365,7 +366,29 @@ def readline_expect(
"""
if timeout is None:
timeout = self.current_default_timeout
return readline_expect(self, expected, read_until, timeout)

expected = ensure_bytes(expected)
read_until = ensure_bytes(read_until)

read = self.read_until(read_until, timeout)
if read == b"":
err = "Expected {!r} but got nothing/timeout!".format(expected)
self.logger.error(err, stack_info=True)
self.close()
if exception_message:
raise BrokenServiceException(exception_message)
else:
raise BrokenServiceException("Service returned nothing (timeout?).")

if expected not in read:
err = "Expected {!r} but got {!r}".format(expected, read)
self.logger.error(err, stack_info=True)
self.close()
if exception_message:
raise BrokenServiceException(exception_message)
else:
raise BrokenServiceException("Service returned unexpected response.")
return read

def expect(
self,
Expand Down
3 changes: 1 addition & 2 deletions tests/test_enochecker.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
ensure_bytes,
ensure_unicode,
parse_args,
readline_expect,
run,
serve_once,
snake_caseify,
Expand Down Expand Up @@ -248,7 +247,7 @@ def test_checker_connections():
checker = CheckerExampleImpl(CHECKER_METHODS[0])
t = checker.connect()
t.write(b"GET / HTTP/1.0\r\n\r\n")
assert readline_expect(t, "HTTP")
assert t.readline_expect("HTTP")
t.close()


Expand Down

0 comments on commit 864d4ae

Please sign in to comment.