distributed.camp
A Python port of the Invisible Internet Project (I2P)
1"""NTCP2 transport protocol types."""
2
3import enum
4import struct
5
6
7class FrameType(enum.IntEnum):
8 DATA = 0
9 PADDING = 254
10 DATETIME = 1
11 OPTIONS = 2
12 ROUTER_INFO = 3
13 I2NP = 4
14 TERMINATION = 5
15
16
17class TransportState(enum.Enum):
18 UNKNOWN = 0
19 HANDSHAKING = 1
20 ESTABLISHED = 2
21 TERMINATED = 3
22 ERROR = 4
23
24
25class NTCP2Frame:
26 """NTCP2 block/frame: type(1) + length(2) + payload."""
27
28 def __init__(self, frame_type: FrameType, payload: bytes):
29 self.frame_type = frame_type
30 self.payload = payload
31
32 def to_bytes(self) -> bytes:
33 return struct.pack("!BH", self.frame_type, len(self.payload)) + self.payload
34
35 @classmethod
36 def from_bytes(cls, data: bytes) -> "NTCP2Frame":
37 frame_type, length = struct.unpack("!BH", data[:3])
38 return cls(FrameType(frame_type), data[3:3 + length])
39
40 @classmethod
41 def from_stream(cls, stream) -> "NTCP2Frame":
42 header = stream.read(3)
43 frame_type, length = struct.unpack("!BH", header)
44 payload = stream.read(length)
45 return cls(FrameType(frame_type), payload)
46
47
48class NTCP2SessionState:
49 """NTCP2 session state machine."""
50
51 def __init__(self):
52 self.state = TransportState.UNKNOWN
53 self.error_message: str | None = None
54 self.messages_sent = 0
55 self.messages_received = 0
56 self._termination_reason: int | None = None
57
58 def begin_handshake(self):
59 if self.state != TransportState.UNKNOWN:
60 raise RuntimeError(f"Cannot begin handshake from {self.state}")
61 self.state = TransportState.HANDSHAKING
62
63 def complete_handshake(self):
64 if self.state != TransportState.HANDSHAKING:
65 raise RuntimeError(f"Cannot complete handshake from {self.state}")
66 self.state = TransportState.ESTABLISHED
67
68 def terminate(self, reason: int = 0):
69 self._termination_reason = reason
70 self.state = TransportState.TERMINATED
71
72 def set_error(self, message: str):
73 self.error_message = message
74 self.state = TransportState.ERROR
75
76 def on_message_sent(self):
77 self.messages_sent += 1
78
79 def on_message_received(self):
80 self.messages_received += 1