From f2d7801135e21cc5de36203f8b39ae98be9d03dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Ricks?= Date: Tue, 18 Jun 2024 09:37:36 +0200 Subject: [PATCH] Adjust transform classes for bytes based protocol For the protocol implementation bytes are passed in the form of Responses to the transform callables. Therefore adjust the Transform classes to handle bytes now. --- gvm/transforms.py | 12 +++++++----- tests/transforms/test_check_command_transform.py | 12 ++++++------ .../transforms/test_etree_check_command_transform.py | 6 +++--- tests/transforms/test_etree_transform.py | 4 ++-- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/gvm/transforms.py b/gvm/transforms.py index 7add892da..5f4b42edf 100644 --- a/gvm/transforms.py +++ b/gvm/transforms.py @@ -7,6 +7,8 @@ """ from lxml import etree +from gvm.protocols.core import Response + from .errors import GvmError, GvmResponseError, GvmServerError from .xml import Element, create_parser @@ -19,10 +21,10 @@ class EtreeTransform: def __init__(self): self._parser = create_parser() - def _convert_response(self, response: str) -> Element: - return etree.XML(response, parser=self._parser) + def _convert_response(self, response: Response) -> Element: + return etree.XML(bytes(response), parser=self._parser) - def __call__(self, response: str) -> Element: + def __call__(self, response: Response) -> Element: return self._convert_response(response) @@ -46,7 +48,7 @@ class CheckCommandTransform(EtreeTransform): response was an error response """ - def __call__(self, response: str) -> str: # type: ignore[override] + def __call__(self, response: Response) -> Response: # type: ignore[override] root = self._convert_response(response) check_command_status(root) @@ -60,7 +62,7 @@ class EtreeCheckCommandTransform(EtreeTransform): response was an error response """ - def __call__(self, response: str) -> Element: + def __call__(self, response: Response) -> Element: root = self._convert_response(response) check_command_status(root) diff --git a/tests/transforms/test_check_command_transform.py b/tests/transforms/test_check_command_transform.py index 15615a8e5..0ea4ddb4c 100644 --- a/tests/transforms/test_check_command_transform.py +++ b/tests/transforms/test_check_command_transform.py @@ -16,7 +16,7 @@ def test_no_status_transform(self): transform = CheckCommandTransform() with self.assertRaises(GvmError): - transform("") + transform(b"") def test_no_success_300status_transform(self): transform = CheckCommandTransform() @@ -25,7 +25,7 @@ def test_no_success_300status_transform(self): root.set("status", "300") root.set("status_text", "Foo error") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) with self.assertRaises(GvmError): transform(response) @@ -37,7 +37,7 @@ def test_no_success_400status_transform(self): root.set("status", "400") root.set("status_text", "Foo error") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) with self.assertRaises(GvmResponseError): transform(response) @@ -49,7 +49,7 @@ def test_no_success_500status_transform(self): root.set("status", "500") root.set("status_text", "Foo error") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) with self.assertRaises(GvmServerError): transform(response) @@ -60,6 +60,6 @@ def test_success_response(self): root = etree.Element("foo_response") root.set("status", "200") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) - self.assertEqual(transform(response), '') + self.assertEqual(transform(response), b'') diff --git a/tests/transforms/test_etree_check_command_transform.py b/tests/transforms/test_etree_check_command_transform.py index 61e04e704..38e0b0deb 100644 --- a/tests/transforms/test_etree_check_command_transform.py +++ b/tests/transforms/test_etree_check_command_transform.py @@ -16,7 +16,7 @@ def test_no_status_transform(self): transform = EtreeCheckCommandTransform() with self.assertRaises(GvmError): - transform("") + transform(b"") def test_no_success_status_transform(self): transform = EtreeCheckCommandTransform() @@ -25,7 +25,7 @@ def test_no_success_status_transform(self): root.set("status", "400") root.set("status_text", "Foo error") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) with self.assertRaises(GvmError): transform(response) @@ -36,7 +36,7 @@ def test_success_response(self): root = etree.Element("foo_response") root.set("status", "200") - response = etree.tostring(root).decode("utf-8") + response = etree.tostring(root) result = transform(response) diff --git a/tests/transforms/test_etree_transform.py b/tests/transforms/test_etree_transform.py index feeb11e56..b4dfe08e0 100644 --- a/tests/transforms/test_etree_transform.py +++ b/tests/transforms/test_etree_transform.py @@ -13,13 +13,13 @@ class EtreeTransformTestCase(unittest.TestCase): def test_transform_response(self): transform = EtreeTransform() - result = transform("") + result = transform(b"") self.assertTrue(etree.iselement(result)) def test_transform_more_complex_response(self): transform = EtreeTransform() - result = transform('') + result = transform(b'') self.assertTrue(etree.iselement(result)) self.assertEqual(result.tag, "foo")