HackTheBox Space Heist Challenge

private static final String SECRET = "qvb4a1b07E870B";

#!/usr/bin/python3
from typing import List, Tuple, Dict
from dataclasses import dataclass
import socket,struct,hashlib,sys
import matplotlib.pyplot as plt
from time import sleep

@dataclass
class TraceData:
    label: str
    samples: List[float]

class PowerAnalyzer:
    def __init__(self, host: str, port: int, plot_enabled: bool = False):
        self.host = host
        self.port = port
        self.plot_enabled = plot_enabled
        self.valid_chars = "0123456789"
        
    def _recv_bytes(self, conn: socket.socket, n_bytes: int) -> bytes:
        buffer = b''
        while len(buffer) < n_bytes:
            buffer += conn.recv(n_bytes - len(buffer))
        return buffer
    
    def _socket_readline(self, conn: socket.socket) -> bytes:
        buffer = b''
        while True:
            char = self._recv_bytes(conn, 1)
            buffer += char
            if char == b'\n':
                break
        return buffer

    def _recv_trace(self, conn: socket.socket) -> TraceData:
        label_len = struct.unpack("<L", self._recv_bytes(conn, 4))[0]
        label = self._recv_bytes(conn, label_len).decode()
        trace_len = struct.unpack("<L", self._recv_bytes(conn, 4))[0]
        sample_buffer = self._recv_bytes(conn, trace_len * 4)
        samples = [struct.unpack("<f", sample_buffer[i:i+4])[0] 
                  for i in range(0, len(sample_buffer), 4)]
        return TraceData(label, samples)

    def _find_peaks(self, trace: List[float]) -> List[Tuple[int, float]]:
        min_val = min(trace)
        max_val = max(trace)
        threshold = (max_val - min_val) / 2 + min_val
        peaks = []
        state = 0
        for i, value in enumerate(trace):
            if state == 0 and value > threshold:
                state = 1
                peaks.append((i, value))
            elif state == 1 and value < threshold:
                state = 0
        return peaks

    def _calc_delta_peaks(self, peaks: List[Tuple[int, float]]) -> List[int]:
        return [curr[0] - prev[0] for prev, curr in zip([(0, 0)] + peaks[:-1], peaks)]

    def _plot_traces(self, traces: Dict[str, List[float]]):
        plt.figure(figsize=(12, 6))
        colors = {'trace_led_auth': 'r', 'trace_led_unlocked': 'g', 'trace_mcu': 'b'}
        for label, trace in traces.items():
            plt.plot(trace, colors.get(label, 'k'), label=label)
        plt.legend()
        plt.grid(True)
        plt.show(block=True)

    def _run_attempt(self, conn: socket.socket, code: str) -> Dict[str, List[float]]:
        print(f'Attempting code: {code}')
        conn.send(code.encode() + b'\n')
        traces = {}
        for _ in range(3):
            trace_data = self._recv_trace(conn)
            traces[trace_data.label] = trace_data.samples
        return traces

    def analyze_power_trace(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as conn:
            conn.connect((self.host, self.port))
            print(self._socket_readline(conn).decode())
            print(self._socket_readline(conn).decode())
            conn.recv(1024)
            conn.send(b'password')
            conn.recv(1024)
            code = ['A']
            last_delta = None
            char_pos = 0
            curr_digit = 0
            while True:
                if curr_digit >= len(self.valid_chars):
                    print(f'Failed to determine char at position {char_pos}')
                    break 
                code[char_pos] = self.valid_chars[curr_digit]
                curr_digit += 1
                traces = self._run_attempt(conn, ''.join(code))
                if self.plot_enabled:
                    self._plot_traces(traces)
                self._socket_readline(conn)
                peaks = self._find_peaks(traces['trace_mcu'])
                delta_peaks = self._calc_delta_peaks(peaks)
                if last_delta is None:
                    last_delta = delta_peaks[-1]
                    continue
                if last_delta - delta_peaks[-1] > 50:
                    min_val = min(traces['trace_led_unlocked'])
                    max_val = max(traces['trace_led_unlocked'])
                    if max_val - min_val > 1:
                        print('Successfully unlocked!')
                        break  
                    last_delta = delta_peaks[-1]
                    char_pos += 1
                    curr_digit = 0
                    code.append('A')
            if self.plot_enabled:
                self._plot_traces(traces)
            return ''.join(code)

    def authenticate(self, code: str, secret: str = "qvb4a1b07E870B"):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as conn:
            conn.connect((self.host, self.port))
            for _ in range(2):
                print(conn.recv(1024).strip().decode())
            conn.send(b'auth')
            for _ in range(2):
                print(conn.recv(1024).strip().decode())
            sleep(0.1)
            conn.send(b'getChallenge')
            response = conn.recv(1024).strip().decode()
            challenge = response.split(":")[1]
            message = challenge + secret
            sha_hash = hashlib.sha1(message.encode()).hexdigest()
            response = f'resp:{challenge}:{sha_hash}'.encode()
            conn.send(response)
            for _ in range(2):
                print(conn.recv(1024).strip().decode())
            conn.send(code.encode())
            for _ in range(2):
                print(conn.recv(1024).strip().decode())

def main():
    if len(sys.argv) < 2:
        print("Usage: python script.py <ip:port> [plt]")
        sys.exit(1) 
    host, port = sys.argv[1].split(":")
    plot_enabled = len(sys.argv) == 3 and sys.argv[2].lower() == 'plt'
    analyzer = PowerAnalyzer(host, int(port), plot_enabled)
    code = analyzer.analyze_power_trace()
    analyzer.authenticate(code)

if __name__ == "__main__":
    main()