test: add tests for x509.Peer

contain hello0, hello, decode, encode

re6st/tests/test_unit/test_x509.py

+import unittest
+import tempfile
+import os
+import struct
+import hashlib
+from mock import Mock, patch
+
+
+from OpenSSL import crypto
+from re6st.x509 import Peer, PACKED_PROTOCOL, fingerprint
+from re6st.tests import tools
+
+
+PROTOCOL = 7
+
+class TestPeer(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.ca_key, cls.ca = tools.create_ca_file(os.devnull, os.devnull)
+        cls.serial = 16
+
+    def setUp(self):
+        self.p = Peer(self.prefix)
+
+    @property
+    def prefix(self):
+        return format(self.serial, '016b')
+
+    @property
+    def cert(self):
+        self._key, self._cert = tools.create_cert_file(os.devnull, os.devnull,
+                                                       self.ca, self.ca_key,
+                                                       self.prefix, self.serial)
+
+        return crypto.load_certificate(crypto.FILETYPE_PEM, self._cert)
+
+
+    def test_gt(self):
+        p2 = Peer(format(0, '016b'))
+
+        self.assertGreater(self.p, p2)
+
+    def test_lt(self):
+        p2 = format((1<<17)-1, '016b')
+
+        self.assertLess(self.p, p2)
+
+    @patch("OpenSSL.crypto.dump_certificate")
+    def test_hello0_0(self, crypto):
+        "0, A"
+        cert = "a cert"
+        crypto.return_value = b"cert"
+
+        res = self.p.hello0(cert)
+
+        self.assertEqual(b'\0\0\0\0' + b"cert", res)
+
+    def test_hello0_1(self):
+        "1, protocol, fingerprint(B), A"
+        self.p.cert = self.cert
+        self.serial += 1
+        cert = self.cert
+
+        res = self.p.hello0(cert)
+
+        expect = (b'\0\0\0\1'
+                 + PACKED_PROTOCOL
+                 + fingerprint(self.p.cert).digest()
+                 + crypto.dump_certificate(crypto.FILETYPE_ASN1, cert))
+        self.assertEqual(expect, res)
+
+    def test_hello(self):
+        self.p.cert = self.cert
+        cert = Mock()
+        sign = cert.sign.return_value = b"a sign"
+
+        res = self.p.hello(cert, PROTOCOL)
+
+        self.assertTrue(res.startswith(b'\0\0\0\2' + PACKED_PROTOCOL))
+        self.assertTrue(res.endswith(sign))
+        # verify encrypt key
+        code = res[len(b'\0\0\0\2' + PACKED_PROTOCOL): -len(sign)]
+        decode_key = tools.decrypt(self._key, code)
+        self.assertEqual(self.p._key, decode_key)
+        self.assertEqual(self.p._i, 2)
+        self.assertEqual(self.p._j, 2)
+
+    def test_decode_valid(self):
+        """seqno > 2"""
+        self.p._i = 2
+        seqno = 3
+        seqno = struct.pack("!L", seqno)
+        msg = b"a msg"
+        _hmac = self.p._hmac(seqno + msg)
+
+        res = self.p.decode(seqno + msg + _hmac)
+
+        self.assertEqual(res, msg)
+
+    def test_decode_wrong_hamc(self):
+        """seqno > 2 but wrong hmac"""
+        self.p._i = 2
+        seqno = 3
+        seqno = struct.pack("!L", seqno)
+        msg = b"a msg"
+        _hmac = "self.p._hmac(seqno + msg)"
+
+        res = self.p.decode(seqno + msg + _hmac)
+
+        self.assertEqual(res, None)
+
+    def test_decode_0(self):
+        """seqno = 0"""
+        seqno = b'\0\0\0\0'
+        msg = b"a msg"
+
+        res = self.p.decode(seqno + msg)
+
+        self.assertEqual(res, (0, msg, None))
+
+    def test_decode_1(self):
+        """seqno = 1"""
+        seqno = b'\0\0\0\1'
+        msg = b"a msg"
+
+        res = self.p.decode(seqno + PACKED_PROTOCOL +msg)
+
+        self.assertEqual(res, (1, msg, PROTOCOL))
+
+    @patch("re6st.x509.Peer._hmac")
+    def test_encode(self, _hmac):
+        new_j = self.p._j = 2
+        new_j += 1
+        msg = b"a msg"
+        hmac_msg = _hmac.return_value = b'a hmac'
+
+        res = self.p.encode(b"a msg")
+
+        self.assertEqual(struct.pack("!L", new_j) + msg + hmac_msg, res)
+        self.assertEqual(new_j, self.p._j)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file