feat: add protocol data encrypt and decrypt module (#8)
diff --git a/Makefile b/Makefile
index 4e88c6c..8ad5cf4 100644
--- a/Makefile
+++ b/Makefile
@@ -20,3 +20,8 @@
python3 -m pip install a6pluginprotos --ignore-installed
python3 -m pip install minicache --ignore-installed
python3 -m pip install click --ignore-installed
+ python3 -m pip install pytest --ignore-installed
+
+.PHONY: test
+test:
+ pytest .
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..d207ff6
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,20 @@
+;
+; Licensed to the Apache Software Foundation (ASF) under one or more
+; contributor license agreements. See the NOTICE file distributed with
+; this work for additional information regarding copyright ownership.
+; The ASF licenses this file to You under the Apache License, Version 2.0
+; (the "License"); you may not use this file except in compliance with
+; the License. You may obtain a copy of the License at
+;
+; http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+;
+[pytest]
+addopts = -vs -p no:warnings
+testpaths = tests
+python_files = test_*
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..b1312a0
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1,16 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
diff --git a/src/runner/socket/protocol.py b/src/runner/socket/protocol.py
new file mode 100644
index 0000000..e412981
--- /dev/null
+++ b/src/runner/socket/protocol.py
@@ -0,0 +1,86 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from src.runner.socket.response import Response as NewServerResponse
+from src.runner.socket.response import RUNNER_ERROR_CODE
+from src.runner.socket.response import RUNNER_SUCCESS_CODE
+from src.runner.socket.response import RUNNER_SUCCESS_MESSAGE
+
+
+class Protocol:
+
+ def __init__(self, buffer: bytes = b'', ty: int = 0):
+ self.__buffer = buffer
+ self.__type = ty
+ self.__length = 0
+
+ @property
+ def length(self) -> int:
+ """
+ get buffer length
+ :return:
+ """
+ return self.__length
+
+ @property
+ def type(self) -> int:
+ """
+ get protocol type
+ :return:
+ """
+ return self.__type
+
+ @property
+ def buffer(self) -> bytes:
+ """
+ get buffer data
+ :return:
+ """
+ return self.__buffer
+
+ def encode(self) -> NewServerResponse:
+ """
+ encode protocol buffer data
+ :return:
+ """
+ if len(self.__buffer) == 0:
+ return NewServerResponse(RUNNER_ERROR_CODE, "ERR: send buffer is empty")
+ response_len = len(self.__buffer)
+ response_header = response_len.to_bytes(4, byteorder="big")
+ response_header = bytearray(response_header)
+ response_header[0] = self.__type
+ response_header = bytes(response_header)
+ self.__buffer = response_header + self.__buffer
+ self.__length = len(self.__buffer)
+ return NewServerResponse(code=RUNNER_SUCCESS_CODE, message=RUNNER_SUCCESS_MESSAGE)
+
+ def decode(self) -> NewServerResponse:
+ """
+ decode protocol buffer data
+ :return:
+ """
+ if len(self.__buffer) == 0:
+ return NewServerResponse(RUNNER_ERROR_CODE, "ERR: recv buffer is empty")
+ length = len(self.__buffer)
+ if length != 4:
+ return NewServerResponse(RUNNER_ERROR_CODE,
+ "ERR: recv protocol type length is 4, got %d" % length)
+
+ buf = bytearray(self.__buffer)
+ self.__type = buf[0]
+ buf[0] = 0
+ self.__length = int.from_bytes(buf, byteorder="big")
+ return NewServerResponse(code=RUNNER_SUCCESS_CODE, message=RUNNER_SUCCESS_MESSAGE)
diff --git a/src/runner/socket/response.py b/src/runner/socket/response.py
new file mode 100644
index 0000000..08b0f53
--- /dev/null
+++ b/src/runner/socket/response.py
@@ -0,0 +1,59 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from __future__ import annotations
+from a6pluginproto.Err import Code as A6ErrCode
+
+RUNNER_SUCCESS_CODE = 200
+RUNNER_SUCCESS_MESSAGE = "OK"
+RUNNER_ERROR_CODE = 500
+RUNNER_ERROR_MESSAGE = "ERR"
+
+errorCodes = [
+ A6ErrCode.Code.CONF_TOKEN_NOT_FOUND,
+ A6ErrCode.Code.BAD_REQUEST,
+ A6ErrCode.Code.SERVICE_UNAVAILABLE,
+]
+
+
+class Response:
+ def __init__(self, code: int = 0, message: str = '', data: bytes = b'', ty: int = 0):
+ self.__code = code
+ self.__message = message
+ self.__type = ty
+ self.__data = data
+
+ def __eq__(self, other: Response) -> bool:
+ return self.code == other.code and \
+ self.message == other.message and \
+ self.data == other.data and \
+ self.type == other.type
+
+ @property
+ def code(self) -> int:
+ return self.__code
+
+ @property
+ def message(self) -> str:
+ return self.__message
+
+ @property
+ def data(self) -> bytes:
+ return self.__data
+
+ @property
+ def type(self) -> int:
+ return self.__type
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..8e25b19
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,4 @@
+import os
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
diff --git a/tests/runner/socket/test_protocol.py b/tests/runner/socket/test_protocol.py
new file mode 100644
index 0000000..7c1a3c6
--- /dev/null
+++ b/tests/runner/socket/test_protocol.py
@@ -0,0 +1,34 @@
+from src.runner.socket.protocol import Protocol as NewServerProtocol
+from src.runner.http.protocol import RPC_PREPARE_CONF
+from src.runner.socket.response import RUNNER_SUCCESS_CODE
+from src.runner.socket.response import RUNNER_SUCCESS_MESSAGE
+
+
+def test_protocol_encode():
+ buf_str = "Hello Python Runner".encode()
+ protocol = NewServerProtocol(buffer=buf_str, ty=RPC_PREPARE_CONF)
+ err = protocol.encode()
+ buf_len = len(buf_str)
+ buf_arr = bytearray(buf_len.to_bytes(4, byteorder="big"))
+ buf_arr[0] = RPC_PREPARE_CONF
+ buf_data = bytes(buf_arr) + buf_str
+ buf_len = len(buf_data)
+ assert err.code == RUNNER_SUCCESS_CODE
+ assert err.message == RUNNER_SUCCESS_MESSAGE
+ assert protocol.type == RPC_PREPARE_CONF
+ assert protocol.buffer == buf_data
+ assert protocol.length == buf_len
+
+
+def test_protocol_decode():
+ buf_str = "Hello Python Runner".encode()
+ buf_len = len(buf_str)
+ buf_arr = bytearray(buf_len.to_bytes(4, byteorder="big"))
+ buf_arr[0] = RPC_PREPARE_CONF
+ buf_data = bytes(buf_arr)
+ protocol = NewServerProtocol(buffer=buf_data)
+ err = protocol.decode()
+ assert err.code == RUNNER_SUCCESS_CODE
+ assert err.message == RUNNER_SUCCESS_MESSAGE
+ assert protocol.type == RPC_PREPARE_CONF
+ assert protocol.length == buf_len
diff --git a/tests/runner/socket/test_response.py b/tests/runner/socket/test_response.py
new file mode 100644
index 0000000..dc9e88f
--- /dev/null
+++ b/tests/runner/socket/test_response.py
@@ -0,0 +1,43 @@
+from src.runner.socket.response import Response as NewServerResponse
+from src.runner.socket.response import RUNNER_ERROR_CODE
+from src.runner.socket.response import RUNNER_SUCCESS_CODE
+from src.runner.http.protocol import RPC_PREPARE_CONF
+from src.runner.http.protocol import RPC_HTTP_REQ_CALL
+from src.runner.http.protocol import RPC_UNKNOWN
+
+
+def test_response_code():
+ response = NewServerResponse(code=RUNNER_SUCCESS_CODE)
+ assert response.code == RUNNER_SUCCESS_CODE
+ error = NewServerResponse(code=RUNNER_ERROR_CODE)
+ assert error.code == RUNNER_ERROR_CODE
+
+
+def test_response_message():
+ response = NewServerResponse(message="Hello Python Runner")
+ assert response.message == "Hello Python Runner"
+
+
+def test_response_data():
+ response = NewServerResponse(data="Hello Python Runner".encode())
+ assert response.data == b'Hello Python Runner'
+
+
+def test_response_type():
+ response = NewServerResponse(ty=RPC_UNKNOWN)
+ assert response.type == RPC_UNKNOWN
+ response = NewServerResponse(ty=RPC_PREPARE_CONF)
+ assert response.type == RPC_PREPARE_CONF
+ response = NewServerResponse(ty=RPC_HTTP_REQ_CALL)
+ assert response.type == RPC_HTTP_REQ_CALL
+
+
+def test_response_eq():
+ resp1 = NewServerResponse(code=RUNNER_SUCCESS_CODE, message="Hello Python Runner",
+ data="Hello Python Runner".encode(), ty=RPC_PREPARE_CONF)
+ resp2 = NewServerResponse(code=RUNNER_ERROR_CODE, message="Hello Python Runner",
+ data="Hello Python Runner".encode(), ty=RPC_PREPARE_CONF)
+ resp3 = NewServerResponse(code=RUNNER_SUCCESS_CODE, message="Hello Python Runner",
+ data="Hello Python Runner".encode(), ty=RPC_PREPARE_CONF)
+ assert resp1 != resp2
+ assert resp1 == resp3
