Merge ebbfbf1 into bad2b98

tests/tlstest.py

@@ -857,6 +857,21 @@ def connect():
 
     test_no += 1
 
+    print("Test {0} - Heartbeat extension response callback".format(test_no))
+    synchro.recv(1)
+    connection = connect()
+    settings = HandshakeSettings()
+    heartbeat_payload = os.urandom(100)
+    def heartbeat_response_check(message):
+        assert heartbeat_payload == message.payload
+    settings.heartbeat_response_callback = heartbeat_response_check
+    connection.handshakeClientCert(serverName=address[0], settings=settings)
+    connection.send_heartbeat_request(heartbeat_payload, 16)
+    testConnClient(connection)
+    connection.close()
+
+    test_no += 1
+
     print('Test {0} - good standard XMLRPC https client'.format(test_no))
     address = address[0], address[1]+1
     synchro.recv(1)
@@ -1638,6 +1653,22 @@ def server_bind(self):
 
     test_no += 1
 
+    print("Test {0} - Heartbeat extension response callback".format(test_no))
+    synchro.send(b'R')
+    connection = connect()
+    settings = HandshakeSettings()
+    heartbeat_payload = os.urandom(100)
+    def heartbeat_response_check(message):
+        assert heartbeat_payload == message.payload
+    settings.heartbeat_response_callback = heartbeat_response_check
+    connection.handshakeServer(certChain=x509Chain, privateKey=x509Key,
+                               settings=settings)
+    connection.send_heartbeat_request(heartbeat_payload, 16)
+    testConnServer(connection)
+    connection.close()
+
+    test_no += 1
+
     print("Tests {0}-{1} - XMLRPXC server".format(test_no, test_no + 2))
     test_no += 2
 

tlslite/constants.py

@@ -124,7 +124,8 @@ class ContentType(TLSEnum):
     alert = 21
     handshake = 22
     application_data = 23
-    all = (20, 21, 22, 23)
+    HEARTBEAT = 24  # RFC 6520
+    all = (20, 21, 22, 23, 24)
 
     @classmethod
     def toRepr(cls, value, blacklist=None):
@@ -145,6 +146,7 @@ class ExtensionType(TLSEnum):
     ec_point_formats = 11  # RFC 4492
     srp = 12  # RFC 5054
     signature_algorithms = 13  # RFC 5246
+    HEARTBEAT = 15  # RFC 6520
     alpn = 16  # RFC 7301
     client_hello_padding = 21  # RFC 7685
     encrypt_then_mac = 22  # RFC 7366
@@ -365,6 +367,20 @@ class CertificateStatusType(TLSEnum):
     ocsp = 1
 
 
+class HeartbeatMode(TLSEnum):
+    """Types of heartbeat modes from RFC 6520"""
+
+    PEER_ALLOWED_TO_SEND = 1
+    PEER_NOT_ALLOWED_TO_SEND = 2
+
+
+class HeartbeatMessageType(TLSEnum):
+    """Types of heartbeat messages from RFC 6520"""
+
+    HEARTBEAT_REQUEST = 1
+    HEARTBEAT_RESPONSE = 2
+
+
 class AlertLevel(TLSEnum):
     """Enumeration of TLS Alert protocol levels"""
 

tlslite/extensions.py

@@ -1754,6 +1754,45 @@ def write(self, writer):
         writer.bytes += self.key_exchange
 
 
+class HeartbeatExtension(TLSExtension):
+    """
+    Heartbeat extension from RFC 6520
+
+    @type mode: int
+    @ivar mode: mode if peer is allowed or nor allowed to send responses
+    """
+    def __init__(self):
+        super(HeartbeatExtension, self).__init__(
+                                        extType=ExtensionType.HEARTBEAT)
+        self.mode = None
+
+    @property
+    def extData(self):
+        """
+        Return encoded heartbeat mode
+
+        @rtype: bytearray
+        """
+        if self.mode is None:
+            return bytearray(0)
+
+        writer = Writer()
+        writer.add(self.mode, 1)
+
+        return writer.bytes
+
+    def create(self, mode):
+        self.mode = mode
+        return self
+
+    def parse(self, p):
+        self.mode = p.get(1)
+        if p.getRemainingLength() > 0:
+            raise SyntaxError()
+
+        return self
+
+
 class ClientKeyShareExtension(TLSExtension):
     """
     Class for handling the Client Hello version of the Key Share extension.
@@ -2073,6 +2112,7 @@ def __init__(self):
         ExtensionType.supports_npn: NPNExtension,
         ExtensionType.client_hello_padding: PaddingExtension,
         ExtensionType.renegotiation_info: RenegotiationInfoExtension,
+        ExtensionType.HEARTBEAT: HeartbeatExtension,
         ExtensionType.supported_versions: SupportedVersionsExtension,
         ExtensionType.key_share: ClientKeyShareExtension,
         ExtensionType.signature_algorithms_cert:

tlslite/handshakesettings.py

@@ -204,6 +204,14 @@ class HandshakeSettings(object):
         early_data processing. In other words, how many bytes will the server
         try to process, but ignore, in case the Client Hello includes
         early_data extension.
+
+    :vartype use_heartbeat_extension: bool
+    :ivar use_heartbeat_extension: whether to support heartbeat extension from
+        RFC 6520. True by default.
+
+    :vartype heartbeat_response_callback: func
+    :ivar heartbeat_response_callback: Callback to function when Heartbeat
+        response is received
     """
 
     def _init_key_settings(self):
@@ -219,6 +227,8 @@ def _init_key_settings(self):
         self.defaultCurve = "secp256r1"
         self.keyShares = ["secp256r1", "x25519"]
         self.padding_cb = None
+        self.use_heartbeat_extension = True
+        self.heartbeat_response_callback = None
 
     def _init_misc_extensions(self):
         """Default variables for assorted extensions."""
@@ -399,6 +409,13 @@ def _sanityCheckExtensions(other):
         if other.usePaddingExtension not in (True, False):
             raise ValueError("usePaddingExtension must be True or False")
 
+        if other.use_heartbeat_extension not in (True, False):
+            raise ValueError("use_heartbeat_extension must be True or False")
+
+        if other.heartbeat_response_callback and not other.use_heartbeat_extension:
+            raise ValueError("heartbeat_response_callback requires "
+                             "use_heartbeat_extension")
+
         HandshakeSettings._sanityCheckEMSExtension(other)
 
     @staticmethod
@@ -506,6 +523,8 @@ def _copy_key_settings(self, other):
         other.dhGroups = self.dhGroups
         other.defaultCurve = self.defaultCurve
         other.keyShares = self.keyShares
+        other.use_heartbeat_extension = self.use_heartbeat_extension
+        other.heartbeat_response_callback = self.heartbeat_response_callback
 
     def validate(self):
         """

tlslite/messages.py

@@ -20,6 +20,7 @@
 from .utils.tackwrapper import *
 from .utils.deprecations import deprecated_attrs, deprecated_params
 from .extensions import *
+from .utils.format_output import none_as_unknown
 
 
 class RecordHeader(object):
@@ -193,22 +194,15 @@ def write(self):
         w.add(self.description, 1)
         return w.bytes
 
-    @staticmethod
-    def _noneAsUnknown(text, number):
-        """if text is None or empty, format number as 'unknown(number)'"""
-        if not text:
-            text = "unknown({0})".format(number)
-        return text
-
     @property
     def levelName(self):
-        return self._noneAsUnknown(AlertLevel.toRepr(self.level),
-                                   self.level)
+        return none_as_unknown(AlertLevel.toRepr(self.level),
+                             self.level)
 
     @property
     def descriptionName(self):
-        return self._noneAsUnknown(AlertDescription.toRepr(self.description),
-                                   self.description)
+        return none_as_unknown(AlertDescription.toRepr(self.description),
+                             self.description)
 
     def __str__(self):
         return "Alert, level:{0}, description:{1}".format(self.levelName,
@@ -2133,3 +2127,69 @@ def parse(self, p):
 
     def write(self):
         return self.bytes
+
+
+class Heartbeat(object):
+    """
+    Handling Heartbeat messages from RFC 6520
+
+    @type message_type: int
+    @ivar message_type: type of message (response or request)
+
+    @type payload: bytearray
+    @ivar payload: payload
+
+    @type padding: bytearray
+    @ivar padding: random padding of selected length
+    """
+
+    def __init__(self):
+        self.contentType = ContentType.HEARTBEAT
+        self.message_type = 0
+        self.payload = bytearray(0)
+        self.padding = bytearray(0)
+
+    def create(self, message_type, payload, padding_length):
+        """Create heartbeat request or response with selected parameters"""
+        self.message_type = message_type
+        self.payload = payload
+        self.padding = getRandomBytes(padding_length)
+        return self
+
+    def create_response(self):
+        """Creates heartbeat response based on request."""
+        heartbeat_response = Heartbeat().create(
+            HeartbeatMessageType.HEARTBEAT_RESPONSE,
+            self.payload,
+            16)
+        return heartbeat_response
+
+    def parse(self, p):
+        """
+        Deserialize heartbeat message from parser.
+
+        We are reading only message type and payload, ignoring
+        leftover bytes (padding).
+        """
+        self.message_type = p.get(1)
+        self.payload = p.getVarBytes(2)
+        return self
+
+    def write(self):
+        """Serialise heartbeat message."""
+        w = Writer()
+        w.add(self.message_type, 1)
+        w.add(len(self.payload), 2)
+        w.bytes += self.payload
+        w.bytes += self.padding
+        return w.bytes
+
+    @property
+    def messageType(self):
+        """Format heartbeat message to human readable representation."""
+        return none_as_unknown(HeartbeatMessageType.toRepr(
+                             self.message_type), self.message_type)
+
+    def __str__(self):
+        """Return human readable representation of heartbeat message."""
+        return "heartbeat {0}".format(self.messageType)

tlslite/messagesocket.py

@@ -52,7 +52,8 @@ def __init__(self, sock, defragmenter):
         super(MessageSocket, self).__init__(sock)
 
         self.defragmenter = defragmenter
-        self.unfragmentedDataTypes = tuple((ContentType.application_data, ))
+        self.unfragmentedDataTypes = tuple((ContentType.application_data,
+                                            ContentType.HEARTBEAT))
         self._lastRecordVersion = (0, 0)
 
         self._sendBuffer = bytearray(0)