Coverage for kilter/service/session.py: 96.96%

1# Copyright 2022-2025 Dominik Sekotill <dom.sekotill@kodo.org.uk> 

2# 

3# This Source Code Form is subject to the terms of the Mozilla Public 

4# License, v. 2.0. If a copy of the MPL was not distributed with this 

5# file, You can obtain one at http://mozilla.org/MPL/2.0/. 

6 

7""" 

8Sessions are the kernel of a filter, providing it with an async API to access messages 

9""" 

10 

11from __future__ import annotations (empty)

12 

13from collections.abc import AsyncGenerator (empty)

14from collections.abc import AsyncIterator (empty)

15from collections.abc import Sequence (empty)

16from contextvars import ContextVar (empty)

17from dataclasses import dataclass (empty)

18from enum import Enum (empty)

19from ipaddress import IPv4Address (empty)

20from ipaddress import IPv6Address (empty)

21from pathlib import Path (empty)

22from types import TracebackType (empty)

23from typing import AsyncContextManager (empty)

24from typing import Literal (empty)

25from typing import Protocol (empty)

26from typing import TypeAlias (empty)

27from warnings import warn (empty)

28 

29from typing_extensions import Self (empty)

30 

31from ..protocol.core import EditMessage (empty)

32from ..protocol.core import EventMessage (empty)

33from ..protocol.messages import * (empty)

34from . import util (empty)

35 

36FilterResponse: TypeAlias = Accept | Reject | Discard | ReplyCode (empty)

37 

38 

39class Aborted(BaseException): (empty)

40 """ 

41 An exception for aborting filters on receipt of an Abort message 

42 """ 

43 

44 

45class Filter(Protocol): (empty)

46 """ (empty)

47 Filters are callables that accept a `Session` and return a response 

48 """ 

49 

50 async def __call__(self, session: Session, /) -> FilterResponse: ... # noqa: D102 (empty)

51 

52 

53class Sender(Protocol): (empty)

54 """ (empty)

55 Senders asynchronously handle sending messages with their "send" method 

56 """ 

57 

58 async def send(self, message: EditMessage) -> None: ... # noqa: D102 (empty)

59 

60 

61class Phase(int, Enum): (empty)

62 """ 

63 Session phases indicate what messages to expect and are impacted by received messages 

64 

65 Users should not generally need to use these values, however an understanding of the 

66 state-flow they represent is useful for understanding some error exception 

67 raised by `Session` methods. 

68 """ 

69 

70 INIT = 0 (empty)

71 """ (empty)

72 This phase is the pre-connected phase of a session; this phase will be completed before 

73 users see the session object. 

74 """ 

75 

76 CONNECT = 1 (empty)

77 """ (empty)

78 This phase is the starting phase of a session, during which a HELO/EHLO message may be 

79 awaited with `Session.helo()`. 

80 """ 

81 

82 MAIL = 2 (empty)

83 """ (empty)

84 This phase is entered after HELO/EHLO, during which a MAIL message may be awaited with 

85 `Session.envelope_from()`. The `Session.extension()` method may also be used to get 

86 the raw MAIL command with any extension arguments, or any other extension commands 

87 that the MTA does not support (if the MTA supports passing these commands to 

88 a filter). 

89 """ 

90 

91 ENVELOPE = 3 (empty)

92 """ (empty)

93 This phase is entered after MAIL, during which any RCPT commands may be awaited with 

94 `Session.envelope_recipients()`. The `Session.extension()` method may also be used to 

95 get the raw RCPT command with any extension arguments, or any other extension commands 

96 that the MTA does not support (if the MTA supports passing these commands to 

97 a filter). 

98 """ 

99 

100 HEADERS = 4 (empty)

101 """ (empty)

102 This phase is entered after a DATA command, while message headers are processed. 

103 Headers may be iterated as they arrive, or be collected for later through the 

104 `Session.headers` object. 

105 """ 

106 

107 BODY = 5 (empty)

108 """ (empty)

109 This phase is entered after a message's headers have been processed. The raw message 

110 body may be iterated over in chunks through the `Session.body` object. 

111 """ 

112 

113 POST = 6 (empty)

114 """ (empty)

115 This phase is entered once a message's body has been completed (or skipped). During 

116 this phase the message editing methods of a `Session` object or the `Session.headers` 

117 and `Session.body` objects may be used. 

118 """ 

119 

120 

121@dataclass (empty)

122class Position: (empty)

123 """ 

124 A base class for `Before` and `After`, this class is not intended to be used directly 

125 """ 

126 

127 subject: Header|Literal["start"]|Literal["end"] (empty)

128 

129 

130@dataclass (empty)

131class Before(Position): (empty)

132 """ 

133 Indicates a relative position preceding a subject `Header` in a header list 

134 

135 See `HeadersAccessor.insert`. 

136 """ 

137 

138 subject: Header (empty)

139 

140 

141@dataclass (empty)

142class After(Position): (empty)

143 """ 

144 Indicates a relative position following a subject `Header` in a header list 

145 

146 See `HeadersAccessor.insert`. 

147 """ 

148 

149 subject: Header (empty)

150 

151 

152START = Position("start") (empty)

153""" (empty)

154Indicates the start of a header list, before the first (current) header 

155""" 

156 

157END = Position("end") (empty)

158""" (empty)

159Indicates the end of a header list, after the last (current) header 

160""" 

161 

162 

163class Session: (empty)

164 """ 

165 The kernel of a filter, providing an API for filters to access messages from an MTA 

166 """ 

167 

168 host: str (empty)

169 """ (empty)

170 A hostname from a reverse address lookup performed when a client connects 

171 

172 If no name is found this value defaults to the standard presentation format for 

173 `Session.address` surrounded by "[" and "]", e.g. "[192.0.2.100]" 

174 """ 

175 

176 address: IPv4Address|IPv6Address|Path|None (empty)

177 """ (empty)

178 The address of the connected client, or None if unknown 

179 """ 

180 

181 port: int (empty)

182 """ (empty)

183 The port of the connected client if applicable, or 0 otherwise 

184 """ 

185 

186 macros: dict[str, str] (empty)

187 """ (empty)

188 A mapping of string replacements sent by the MTA 

189 

190 See `smfi_getsymval <https://pythonhosted.org/pymilter/milter_api/smfi_getsymval.html>`_ 

191 from `libmilter` for more information. 

192 

193 Warning: 

194 The current implementation is very naïve and does not behave exactly like 

195 `libmilter`, nor is it very robust. It will definitely change in the future. 

196 """ 

197 

198 headers: HeadersAccessor (empty)

199 """ (empty)

200 A `HeadersAccessor` object for accessing and modifying the message header fields 

201 """ 

202 

203 body: BodyAccessor (empty)

204 """ (empty)

205 A `BodyAccessor` object for accessing and modifying the message body 

206 """ 

207 

208 def __init__( (empty)

209 self, 

210 sender: Sender, 

211 broadcast: util.Broadcast[EventMessage]|None = None, 

212 ): 

213 self.host = "" (empty)

214 self.address = None (empty)

215 self.port = 0 (empty)

216 

217 self.sender = sender (empty)

218 self.broadcast = broadcast or util.Broadcast[EventMessage]() (empty)

219 

220 self.macros = dict[str, str]() (empty)

221 self.headers = HeadersAccessor(self, sender) (empty)

222 self.body = BodyAccessor(self, sender) (empty)

223 

224 # Phase checking is a bit fuzzy as a filter may not request every message, 

225 # so some phases will be skipped; checks should not try to exactly match a phase. 

226 self.phase = Phase.INIT (empty)

227 

228 self._helo: Helo|None = None (empty)

229 

230 async def __aenter__(self) -> Self: (empty)

231 await self.broadcast.__aenter__() (empty)

232 return self (empty)

233 

234 async def __aexit__(self, *_: object) -> None: (empty)

235 await self.broadcast.__aexit__(None, None, None) (empty)

236 # on session close, wake up any remaining deliver() awaitables 

237 await self.broadcast.shutdown_hook() (empty)

238 

239 def _reset(self) -> None: (empty)

240 self.headers = HeadersAccessor(self, self.sender) (empty)

241 self.body = BodyAccessor(self, self.sender) (empty)

242 

243 async def deliver(self, message: EventMessage) -> type[Continue]|type[Skip]: (empty)

244 """ 

245 Deliver a message (or its contents) to a task waiting for it 

246 """ 

247 match message: (empty)

248 case Connect(): (empty)

249 self.host = message.hostname (empty)

250 self.address = message.address (empty)

251 self.port = message.port (empty)

252 async with self.broadcast: (empty)

253 self.phase = Phase.CONNECT (empty)

254 return Continue (empty)

255 case Macro(): (empty)

256 self.macros.update(message.macros) (empty)

257 return Continue # not strictly necessary, but type checker needs something (empty)

258 case Abort(): (empty)

259 async with self.broadcast: (empty)

260 self.phase = Phase.CONNECT (empty)

261 await self.broadcast.abort(Aborted) (empty)

262 self._reset() (empty)

263 return Continue (empty)

264 case Helo(): (empty)

265 phase = Phase.MAIL (empty)

266 case EnvelopeFrom() | EnvelopeRecipient() | Unknown(): (empty)

267 phase = Phase.ENVELOPE (empty)

268 case Data() | Header(): (empty)

269 phase = Phase.HEADERS (empty)

270 case EndOfHeaders() | Body(): (empty)

271 phase = Phase.BODY (empty)

272 case EndOfMessage(): # pragma: no-branch (empty)

273 phase = Phase.POST (empty)

274 async with self.broadcast: (empty)

275 self.phase = phase # phase attribute must be modified in locked context (empty)

276 await self.broadcast.send(message) (empty)

277 return Skip if self.phase == Phase.BODY and self.body.should_skip() else Continue (empty)

278 

279 async def helo(self) -> str: (empty)

280 """ 

281 Wait for a HELO/EHLO message and return the client's claimed hostname 

282 """ 

283 if self.phase > Phase.CONNECT: (empty)

284 raise RuntimeError( (empty)

285 "Session.helo() must be awaited before any other async features of a " 

286 "Session", 

287 ) 

288 if self._helo: (empty)

289 return self._helo.hostname (empty)

290 while self.phase <= Phase.CONNECT: (empty)

291 message = await self.broadcast.receive() (empty)

292 if isinstance(message, Helo): (empty)

293 self._helo = message (empty)

294 return message.hostname (empty)

295 raise RuntimeError("HELO/EHLO event not received") (empty)

296 

297 async def envelope_from(self) -> str: (empty)

298 """ 

299 Wait for a MAIL command message and return the sender identity 

300 

301 Note that if extensions arguments are wanted, users should use `Session.extension()` 

302 instead with a name of ``"MAIL"``. 

303 """ 

304 if self.phase > Phase.MAIL: (empty)

305 raise RuntimeError( (empty)

306 "Session.envelope_from() may only be awaited before the ENVELOPE phase", 

307 ) 

308 while self.phase <= Phase.MAIL: (empty)

309 message = await self.broadcast.receive() (empty)

310 if isinstance(message, EnvelopeFrom): (empty)

311 return bytes(message.sender).decode() (empty)

312 raise RuntimeError("MAIL event not received") (empty)

313 

314 async def envelope_recipients(self) -> AsyncIterator[str]: (empty)

315 """ 

316 Wait for RCPT command messages and iteratively yield the recipients' identities 

317 

318 Note that if extensions arguments are wanted, users should use `Session.extension()` 

319 instead with a name of ``"RCPT"``. 

320 """ 

321 if self.phase > Phase.ENVELOPE: (empty)

322 raise RuntimeError( (empty)

323 "Session.envelope_from() may only be awaited before the HEADERS phase", 

324 ) 

325 while self.phase <= Phase.ENVELOPE: (empty)

326 message = await self.broadcast.receive() (empty)

327 if isinstance(message, EnvelopeRecipient): (empty)

328 yield bytes(message.recipient).decode() (empty)

329 

330 async def extension(self, name: str) -> memoryview: (empty)

331 """ 

332 Wait for the named command extension and return the raw command for processing 

333 """ 

334 if self.phase > Phase.ENVELOPE: (empty)

335 raise RuntimeError( (empty)

336 "Session.extension() may only be awaited before the HEADERS phase", 

337 ) 

338 bname = name.encode("utf-8") (empty)

339 while self.phase <= Phase.ENVELOPE: (empty)

340 message = await self.broadcast.receive() (empty)

341 match message: (empty)

342 case Unknown(): (empty)

343 if message.content[:len(bname)] == bname: (empty)

344 assert isinstance(message.content, memoryview) (empty)

345 return message.content (empty)

346 # fake buffers for MAIL and RCPT commands 

347 case EnvelopeFrom() if name == "MAIL": (empty)

348 vals = [b"MAIL FROM", message.sender, *message.arguments] (empty)

349 return memoryview(b" ".join(vals)) (empty)

350 case EnvelopeRecipient() if name == "RCPT": (empty)

351 vals = [b"RCPT TO", message.recipient, *message.arguments] (empty)

352 return memoryview(b" ".join(vals)) (empty)

353 raise RuntimeError(f"{name} event not received") (empty)

354 

355 async def change_sender(self, sender: str, args: str = "") -> None: (empty)

356 """ 

357 Move onto the `Phase.POST` phase and instruct the MTA to change the sender address 

358 """ 

359 await _until_editable(self) (empty)

360 await self.sender.send(ChangeSender(sender, args or None)) (empty)

361 

362 async def add_recipient(self, recipient: str, args: str = "") -> None: (empty)

363 """ 

364 Move onto the `Phase.POST` phase and instruct the MTA to add a new recipient address 

365 """ 

366 await _until_editable(self) (empty)

367 await self.sender.send( (empty)

368 AddRecipientPar(recipient, args) if args else AddRecipient(recipient), 

369 ) 

370 

371 async def remove_recipient(self, recipient: str) -> None: (empty)

372 """ 

373 Move onto the `Phase.POST` phase and instruct the MTA to remove a recipient address 

374 """ 

375 await _until_editable(self) (empty)

376 await self.sender.send(RemoveRecipient(recipient)) (empty)

377 

378 

379class HeadersAccessor(AsyncContextManager["HeaderIterator"]): (empty)

380 """ 

381 A class that allows access and modification of the message headers sent from an MTA 

382 

383 To access headers (which are only available iteratively), use an instance as an 

384 asynchronous context manager; a `HeaderIterator` is returned when the context is 

385 entered. 

386 """ 

387 

388 def __init__(self, session: Session, sender: Sender): (empty)

389 self.session = session (empty)

390 self.sender = sender (empty)

391 self._table = list[Header]() (empty)

392 self._aiter = ContextVar[HeaderIterator|None]("header-iter") (empty)

393 

394 async def __aenter__(self) -> HeaderIterator: (empty)

395 if not (aiter := self._aiter.get(None)): 395 ↛ 398line 395 didn't jump to line 398 because the condition on line 395 was always true(empty)

396 aiter = HeaderIterator(self.__aiter()) (empty)

397 self._aiter.set(aiter) (empty)

398 return aiter (empty)

399 

400 async def __aexit__(self, *_: object) -> None: (empty)

401 if aiter := self._aiter.get(): 401 ↛ 403line 401 didn't jump to line 403 because the condition on line 401 was always true(empty)

402 await aiter.aclose() (empty)

403 self._aiter.set(None) (empty)

404 

405 async def __aiter(self) -> AsyncGenerator[Header, None]: (empty)

406 # yield from cached headers first; allows multiple tasks to access the headers 

407 # in an uncoordinated manner; note the broadcaster is locked at this point 

408 for header in self._table: (empty)

409 yield header (empty)

410 seen = set(id(header) for header in self._table) (empty)

411 while self.session.phase <= Phase.HEADERS: (empty)

412 match (await self.session.broadcast.receive()): (empty)

413 case Header() as header: (empty)

414 header.freeze() (empty)

415 self._table.append(header) (empty)

416 seen.add(id(header)) (empty)

417 try: (empty)

418 yield header (empty)

419 except GeneratorExit: (empty)

420 await self.collect() (empty)

421 raise (empty)

422 case EndOfHeaders(): (empty)

423 return (empty)

424 # It's possible for collect() to have been called while yielded, in which case the 

425 # loop will end. Yield any headers that were stored by collect() but not yet 

426 # yielded. 

427 for header in self._table: (empty)

428 if id(header) not in seen: 428 ↛ 429line 428 didn't jump to line 429 because the condition on line 428 was never true(empty)

429 yield header 

430 

431 async def collect(self) -> None: (empty)

432 """ 

433 Collect all headers without producing an iterator 

434 

435 Calling this method before the `Phase.BODY` phase allows later processing of headers 

436 (after the HEADER phase) without the need for an empty loop. 

437 """ 

438 # note the similarities between this and __aiter; the difference is no mutex or 

439 # yields 

440 while self.session.phase <= Phase.HEADERS: (empty)

441 match (await self.session.broadcast.receive()): (empty)

442 case Header() as header: (empty)

443 header.freeze() (empty)

444 self._table.append(header) (empty)

445 case _: (empty)

446 return (empty)

447 

448 async def delete(self, header: Header) -> None: (empty)

449 """ 

450 Move onto the `Phase.POST` phase and Instruct the MTA to delete the given header 

451 """ 

452 await self.collect() (empty)

453 await _until_editable(self.session) (empty)

454 index = _index_by_name(self._table, header) (empty)

455 await self.sender.send(ChangeHeader(index, header.name, b"")) (empty)

456 self._table.remove(header) (empty)

457 

458 async def update(self, header: Header, value: bytes) -> None: (empty)

459 """ 

460 Move onto the `Phase.POST` phase and Instruct the MTA to modify the value of a header 

461 """ 

462 await self.collect() (empty)

463 await _until_editable(self.session) (empty)

464 index = _index_by_name(self._table, header) (empty)

465 await self.sender.send(ChangeHeader(index, header.name, value)) (empty)

466 index = self._table.index(header) (empty)

467 self._table[index].value = value (empty)

468 

469 async def insert(self, header: Header, position: Position) -> None: (empty)

470 """ 

471 Move onto the `Phase.POST` phase and instruct the MTA to insert a new header 

472 

473 The header is inserted at `START`, `END`, or a relative position with `Before` and 

474 `After`; for example ``Before(Header("To", "test@example.com"))``. 

475 """ 

476 await self.collect() (empty)

477 await _until_editable(self.session) (empty)

478 match position: (empty)

479 case Position(subject="start"): (empty)

480 index = 0 (empty)

481 case Position(subject="end"): (empty)

482 index = len(self._table) (empty)

483 case Before(): (empty)

484 index = self._table.index(position.subject) (empty)

485 case After(): # pragma: no-branch (empty)

486 index = self._table.index(position.subject) + 1 (empty)

487 case _: 

488 raise TypeError("Expect a Position") 

489 if index >= len(self._table): (empty)

490 await self.sender.send(AddHeader(header.name, header.value)) (empty)

491 self._table.append(header) (empty)

492 else: 

493 await self.sender.send(InsertHeader(index + 1, header.name, header.value)) (empty)

494 self._table.insert(index, header) (empty)

495 

496 

497class HeaderIterator(AsyncGenerator[Header, None]): (empty)

498 """ 

499 Iterator for headers obtained by using a `HeadersAccessor` as a context manager 

500 """ 

501 

502 def __init__(self, aiter: AsyncGenerator[Header, None]): (empty)

503 self._aiter = aiter (empty)

504 

505 def __aiter__(self) -> Self: (empty)

506 return self (empty)

507 

508 async def __anext__(self) -> Header: # noqa: D102 (empty)

509 return await self._aiter.__anext__() (empty)

510 

511 async def asend(self, value: None = None) -> Header: # noqa: D102 (empty)

512 return await self._aiter.__anext__() (empty)

513 

514 async def athrow( # noqa: D102 (empty)

515 self, 

516 e: type[BaseException]|BaseException, 

517 m: object = None, 

518 t: TracebackType|None = None, /, 

519 ) -> Header: 

520 if isinstance(e, type): (empty)

521 return await self._aiter.athrow(e, m, t) (empty)

522 assert m is None (empty)

523 return await self._aiter.athrow(e, m, t) (empty)

524 

525 async def aclose(self) -> None: # noqa: D102 (empty)

526 await self._aiter.aclose() (empty)

527 

528 async def restrict(self, *names: str) -> AsyncIterator[Header]: (empty)

529 """ 

530 Return an asynchronous generator that filters headers by name 

531 """ 

532 async for header in self._aiter: (empty)

533 if header.name in names: (empty)

534 yield header (empty)

535 

536 

537class BodyAccessor(AsyncContextManager[AsyncIterator[memoryview]]): (empty)

538 """ 

539 A class that allows access and modification of the message body sent from an MTA 

540 

541 To access chunks of a body (which are only available iteratively), use an instance as an 

542 asynchronous context manager; an asynchronous iterator is returned when the context is 

543 entered. 

544 """ 

545 

546 def __init__(self, session: Session, sender: Sender): (empty)

547 self.session = session (empty)

548 self.sender = sender (empty)

549 self._entered = 0 (empty)

550 self._skip = False (empty)

551 self._aiter = ContextVar[AsyncGenerator[memoryview, None] | None]("body-iter") (empty)

552 

553 async def __aenter__(self) -> AsyncIterator[memoryview]: (empty)

554 if not (aiter := self._aiter.get(None)): 554 ↛ 557line 554 didn't jump to line 557 because the condition on line 554 was always true(empty)

555 aiter = self.__aiter() (empty)

556 self._aiter.set(aiter) (empty)

557 self._entered += 1 (empty)

558 return aiter (empty)

559 

560 async def __aexit__(self, *_: object) -> None: (empty)

561 if aiter := self._aiter.get(None): 561 ↛ 563line 561 didn't jump to line 563 because the condition on line 561 was always true(empty)

562 await aiter.aclose() (empty)

563 self._aiter.set(None) (empty)

564 self._entered -= 1 (empty)

565 

566 async def __aiter(self) -> AsyncGenerator[memoryview, None]: (empty)

567 while self.session.phase <= Phase.BODY: (empty)

568 match (await self.session.broadcast.receive()): (empty)

569 case Body() as body: (empty)

570 assert isinstance(body.content, memoryview) (empty)

571 yield body.content (empty)

572 case EndOfMessage() as eom: (empty)

573 assert isinstance(eom.content, memoryview) (empty)

574 yield eom.content (empty)

575 

576 def should_skip(self) -> bool: (empty)

577 """ 

578 Return whether the message body should be skipped 

579 

580 The body should be skipped when there are no active contexts. All correctly 

581 implemented filters should have started a context before the first `Body` message. 

582 

583 Once this method returns `True` it becomes "locked in" and will always return `True` 

584 after. 

585 """ 

586 if self._skip: (empty)

587 return True (empty)

588 self._skip = self._entered == 0 (empty)

589 return self._skip (empty)

590 

591 async def write(self, chunk: bytes) -> None: (empty)

592 """ 

593 Request that chunks of a new message body are sent to the MTA 

594 

595 This method should not be called from within the scope created by using it's 

596 instance as an async context (`async with`); doing so may cause a warning to be 

597 issued and the rest of the message body to be skipped. 

598 """ 

599 if self._aiter.get(None): (empty)

600 warn( (empty)

601 "it looks as if BodyAccessor.write() was called on an instance from within " 

602 "it's own async context", 

603 stacklevel=2, 

604 ) 

605 await _until_editable(self.session) (empty)

606 await self.sender.send(ReplaceBody(chunk)) (empty)

607 

608 

609async def _until_editable(session: Session) -> None: (empty)

610 if session.phase == Phase.POST: (empty)

611 return (empty)

612 while session.phase < Phase.POST: (empty)

613 if session.phase == Phase.HEADERS: 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true(empty)

614 await session.headers.collect() 

615 else: 

616 await session.broadcast.receive() (empty)

617 

618 

619def _index_by_name(table: Sequence[Header], needle: Header) -> int: (empty)

620 index = 0 (empty)

621 name = needle.name.lower() (empty)

622 for header in table: 622 ↛ 627line 622 didn't jump to line 627 because the loop on line 622 didn't complete(empty)

623 if header == needle: (empty)

624 return index + 1 (empty)

625 if header.name.lower() == name: 625 ↛ 622line 625 didn't jump to line 622 because the condition on line 625 was always true(empty)

626 index += 1 (empty)

627 raise ValueError(f"header not found: {needle}")