multipart.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957
  1. import base64
  2. import binascii
  3. import json
  4. import re
  5. import uuid
  6. import warnings
  7. import zlib
  8. from collections import deque
  9. from types import TracebackType
  10. from typing import (
  11. TYPE_CHECKING,
  12. Any,
  13. AsyncIterator,
  14. Dict,
  15. Iterator,
  16. List,
  17. Mapping,
  18. Optional,
  19. Sequence,
  20. Tuple,
  21. Type,
  22. Union,
  23. )
  24. from urllib.parse import parse_qsl, unquote, urlencode
  25. from multidict import CIMultiDict, CIMultiDictProxy, MultiMapping
  26. from .hdrs import (
  27. CONTENT_DISPOSITION,
  28. CONTENT_ENCODING,
  29. CONTENT_LENGTH,
  30. CONTENT_TRANSFER_ENCODING,
  31. CONTENT_TYPE,
  32. )
  33. from .helpers import CHAR, TOKEN, parse_mimetype, reify
  34. from .http import HeadersParser
  35. from .payload import (
  36. JsonPayload,
  37. LookupError,
  38. Order,
  39. Payload,
  40. StringPayload,
  41. get_payload,
  42. payload_type,
  43. )
  44. from .streams import StreamReader
  45. __all__ = (
  46. "MultipartReader",
  47. "MultipartWriter",
  48. "BodyPartReader",
  49. "BadContentDispositionHeader",
  50. "BadContentDispositionParam",
  51. "parse_content_disposition",
  52. "content_disposition_filename",
  53. )
  54. if TYPE_CHECKING: # pragma: no cover
  55. from .client_reqrep import ClientResponse
  56. class BadContentDispositionHeader(RuntimeWarning):
  57. pass
  58. class BadContentDispositionParam(RuntimeWarning):
  59. pass
  60. def parse_content_disposition(
  61. header: Optional[str],
  62. ) -> Tuple[Optional[str], Dict[str, str]]:
  63. def is_token(string: str) -> bool:
  64. return bool(string) and TOKEN >= set(string)
  65. def is_quoted(string: str) -> bool:
  66. return string[0] == string[-1] == '"'
  67. def is_rfc5987(string: str) -> bool:
  68. return is_token(string) and string.count("'") == 2
  69. def is_extended_param(string: str) -> bool:
  70. return string.endswith("*")
  71. def is_continuous_param(string: str) -> bool:
  72. pos = string.find("*") + 1
  73. if not pos:
  74. return False
  75. substring = string[pos:-1] if string.endswith("*") else string[pos:]
  76. return substring.isdigit()
  77. def unescape(text: str, *, chars: str = "".join(map(re.escape, CHAR))) -> str:
  78. return re.sub(f"\\\\([{chars}])", "\\1", text)
  79. if not header:
  80. return None, {}
  81. disptype, *parts = header.split(";")
  82. if not is_token(disptype):
  83. warnings.warn(BadContentDispositionHeader(header))
  84. return None, {}
  85. params = {} # type: Dict[str, str]
  86. while parts:
  87. item = parts.pop(0)
  88. if "=" not in item:
  89. warnings.warn(BadContentDispositionHeader(header))
  90. return None, {}
  91. key, value = item.split("=", 1)
  92. key = key.lower().strip()
  93. value = value.lstrip()
  94. if key in params:
  95. warnings.warn(BadContentDispositionHeader(header))
  96. return None, {}
  97. if not is_token(key):
  98. warnings.warn(BadContentDispositionParam(item))
  99. continue
  100. elif is_continuous_param(key):
  101. if is_quoted(value):
  102. value = unescape(value[1:-1])
  103. elif not is_token(value):
  104. warnings.warn(BadContentDispositionParam(item))
  105. continue
  106. elif is_extended_param(key):
  107. if is_rfc5987(value):
  108. encoding, _, value = value.split("'", 2)
  109. encoding = encoding or "utf-8"
  110. else:
  111. warnings.warn(BadContentDispositionParam(item))
  112. continue
  113. try:
  114. value = unquote(value, encoding, "strict")
  115. except UnicodeDecodeError: # pragma: nocover
  116. warnings.warn(BadContentDispositionParam(item))
  117. continue
  118. else:
  119. failed = True
  120. if is_quoted(value):
  121. failed = False
  122. value = unescape(value[1:-1].lstrip("\\/"))
  123. elif is_token(value):
  124. failed = False
  125. elif parts:
  126. # maybe just ; in filename, in any case this is just
  127. # one case fix, for proper fix we need to redesign parser
  128. _value = "{};{}".format(value, parts[0])
  129. if is_quoted(_value):
  130. parts.pop(0)
  131. value = unescape(_value[1:-1].lstrip("\\/"))
  132. failed = False
  133. if failed:
  134. warnings.warn(BadContentDispositionHeader(header))
  135. return None, {}
  136. params[key] = value
  137. return disptype.lower(), params
  138. def content_disposition_filename(
  139. params: Mapping[str, str], name: str = "filename"
  140. ) -> Optional[str]:
  141. name_suf = "%s*" % name
  142. if not params:
  143. return None
  144. elif name_suf in params:
  145. return params[name_suf]
  146. elif name in params:
  147. return params[name]
  148. else:
  149. parts = []
  150. fnparams = sorted(
  151. (key, value) for key, value in params.items() if key.startswith(name_suf)
  152. )
  153. for num, (key, value) in enumerate(fnparams):
  154. _, tail = key.split("*", 1)
  155. if tail.endswith("*"):
  156. tail = tail[:-1]
  157. if tail == str(num):
  158. parts.append(value)
  159. else:
  160. break
  161. if not parts:
  162. return None
  163. value = "".join(parts)
  164. if "'" in value:
  165. encoding, _, value = value.split("'", 2)
  166. encoding = encoding or "utf-8"
  167. return unquote(value, encoding, "strict")
  168. return value
  169. class MultipartResponseWrapper:
  170. """Wrapper around the MultipartReader.
  171. It takes care about
  172. underlying connection and close it when it needs in.
  173. """
  174. def __init__(
  175. self,
  176. resp: "ClientResponse",
  177. stream: "MultipartReader",
  178. ) -> None:
  179. self.resp = resp
  180. self.stream = stream
  181. def __aiter__(self) -> "MultipartResponseWrapper":
  182. return self
  183. async def __anext__(
  184. self,
  185. ) -> Union["MultipartReader", "BodyPartReader"]:
  186. part = await self.next()
  187. if part is None:
  188. raise StopAsyncIteration
  189. return part
  190. def at_eof(self) -> bool:
  191. """Returns True when all response data had been read."""
  192. return self.resp.content.at_eof()
  193. async def next(
  194. self,
  195. ) -> Optional[Union["MultipartReader", "BodyPartReader"]]:
  196. """Emits next multipart reader object."""
  197. item = await self.stream.next()
  198. if self.stream.at_eof():
  199. await self.release()
  200. return item
  201. async def release(self) -> None:
  202. """Releases the connection gracefully, reading all the content
  203. to the void."""
  204. await self.resp.release()
  205. class BodyPartReader:
  206. """Multipart reader for single body part."""
  207. chunk_size = 8192
  208. def __init__(
  209. self, boundary: bytes, headers: "CIMultiDictProxy[str]", content: StreamReader
  210. ) -> None:
  211. self.headers = headers
  212. self._boundary = boundary
  213. self._content = content
  214. self._at_eof = False
  215. length = self.headers.get(CONTENT_LENGTH, None)
  216. self._length = int(length) if length is not None else None
  217. self._read_bytes = 0
  218. # TODO: typeing.Deque is not supported by Python 3.5
  219. self._unread = deque() # type: Any
  220. self._prev_chunk = None # type: Optional[bytes]
  221. self._content_eof = 0
  222. self._cache = {} # type: Dict[str, Any]
  223. def __aiter__(self) -> AsyncIterator["BodyPartReader"]:
  224. return self # type: ignore
  225. async def __anext__(self) -> bytes:
  226. part = await self.next()
  227. if part is None:
  228. raise StopAsyncIteration
  229. return part
  230. async def next(self) -> Optional[bytes]:
  231. item = await self.read()
  232. if not item:
  233. return None
  234. return item
  235. async def read(self, *, decode: bool = False) -> bytes:
  236. """Reads body part data.
  237. decode: Decodes data following by encoding
  238. method from Content-Encoding header. If it missed
  239. data remains untouched
  240. """
  241. if self._at_eof:
  242. return b""
  243. data = bytearray()
  244. while not self._at_eof:
  245. data.extend(await self.read_chunk(self.chunk_size))
  246. if decode:
  247. return self.decode(data)
  248. return data
  249. async def read_chunk(self, size: int = chunk_size) -> bytes:
  250. """Reads body part content chunk of the specified size.
  251. size: chunk size
  252. """
  253. if self._at_eof:
  254. return b""
  255. if self._length:
  256. chunk = await self._read_chunk_from_length(size)
  257. else:
  258. chunk = await self._read_chunk_from_stream(size)
  259. self._read_bytes += len(chunk)
  260. if self._read_bytes == self._length:
  261. self._at_eof = True
  262. if self._at_eof:
  263. clrf = await self._content.readline()
  264. assert (
  265. b"\r\n" == clrf
  266. ), "reader did not read all the data or it is malformed"
  267. return chunk
  268. async def _read_chunk_from_length(self, size: int) -> bytes:
  269. # Reads body part content chunk of the specified size.
  270. # The body part must has Content-Length header with proper value.
  271. assert self._length is not None, "Content-Length required for chunked read"
  272. chunk_size = min(size, self._length - self._read_bytes)
  273. chunk = await self._content.read(chunk_size)
  274. return chunk
  275. async def _read_chunk_from_stream(self, size: int) -> bytes:
  276. # Reads content chunk of body part with unknown length.
  277. # The Content-Length header for body part is not necessary.
  278. assert (
  279. size >= len(self._boundary) + 2
  280. ), "Chunk size must be greater or equal than boundary length + 2"
  281. first_chunk = self._prev_chunk is None
  282. if first_chunk:
  283. self._prev_chunk = await self._content.read(size)
  284. chunk = await self._content.read(size)
  285. self._content_eof += int(self._content.at_eof())
  286. assert self._content_eof < 3, "Reading after EOF"
  287. assert self._prev_chunk is not None
  288. window = self._prev_chunk + chunk
  289. sub = b"\r\n" + self._boundary
  290. if first_chunk:
  291. idx = window.find(sub)
  292. else:
  293. idx = window.find(sub, max(0, len(self._prev_chunk) - len(sub)))
  294. if idx >= 0:
  295. # pushing boundary back to content
  296. with warnings.catch_warnings():
  297. warnings.filterwarnings("ignore", category=DeprecationWarning)
  298. self._content.unread_data(window[idx:])
  299. if size > idx:
  300. self._prev_chunk = self._prev_chunk[:idx]
  301. chunk = window[len(self._prev_chunk) : idx]
  302. if not chunk:
  303. self._at_eof = True
  304. result = self._prev_chunk
  305. self._prev_chunk = chunk
  306. return result
  307. async def readline(self) -> bytes:
  308. """Reads body part by line by line."""
  309. if self._at_eof:
  310. return b""
  311. if self._unread:
  312. line = self._unread.popleft()
  313. else:
  314. line = await self._content.readline()
  315. if line.startswith(self._boundary):
  316. # the very last boundary may not come with \r\n,
  317. # so set single rules for everyone
  318. sline = line.rstrip(b"\r\n")
  319. boundary = self._boundary
  320. last_boundary = self._boundary + b"--"
  321. # ensure that we read exactly the boundary, not something alike
  322. if sline == boundary or sline == last_boundary:
  323. self._at_eof = True
  324. self._unread.append(line)
  325. return b""
  326. else:
  327. next_line = await self._content.readline()
  328. if next_line.startswith(self._boundary):
  329. line = line[:-2] # strip CRLF but only once
  330. self._unread.append(next_line)
  331. return line
  332. async def release(self) -> None:
  333. """Like read(), but reads all the data to the void."""
  334. if self._at_eof:
  335. return
  336. while not self._at_eof:
  337. await self.read_chunk(self.chunk_size)
  338. async def text(self, *, encoding: Optional[str] = None) -> str:
  339. """Like read(), but assumes that body part contains text data."""
  340. data = await self.read(decode=True)
  341. # see https://www.w3.org/TR/html5/forms.html#multipart/form-data-encoding-algorithm # NOQA
  342. # and https://dvcs.w3.org/hg/xhr/raw-file/tip/Overview.html#dom-xmlhttprequest-send # NOQA
  343. encoding = encoding or self.get_charset(default="utf-8")
  344. return data.decode(encoding)
  345. async def json(self, *, encoding: Optional[str] = None) -> Optional[Dict[str, Any]]:
  346. """Like read(), but assumes that body parts contains JSON data."""
  347. data = await self.read(decode=True)
  348. if not data:
  349. return None
  350. encoding = encoding or self.get_charset(default="utf-8")
  351. return json.loads(data.decode(encoding))
  352. async def form(self, *, encoding: Optional[str] = None) -> List[Tuple[str, str]]:
  353. """Like read(), but assumes that body parts contains form
  354. urlencoded data.
  355. """
  356. data = await self.read(decode=True)
  357. if not data:
  358. return []
  359. if encoding is not None:
  360. real_encoding = encoding
  361. else:
  362. real_encoding = self.get_charset(default="utf-8")
  363. return parse_qsl(
  364. data.rstrip().decode(real_encoding),
  365. keep_blank_values=True,
  366. encoding=real_encoding,
  367. )
  368. def at_eof(self) -> bool:
  369. """Returns True if the boundary was reached or False otherwise."""
  370. return self._at_eof
  371. def decode(self, data: bytes) -> bytes:
  372. """Decodes data according the specified Content-Encoding
  373. or Content-Transfer-Encoding headers value.
  374. """
  375. if CONTENT_TRANSFER_ENCODING in self.headers:
  376. data = self._decode_content_transfer(data)
  377. if CONTENT_ENCODING in self.headers:
  378. return self._decode_content(data)
  379. return data
  380. def _decode_content(self, data: bytes) -> bytes:
  381. encoding = self.headers.get(CONTENT_ENCODING, "").lower()
  382. if encoding == "deflate":
  383. return zlib.decompress(data, -zlib.MAX_WBITS)
  384. elif encoding == "gzip":
  385. return zlib.decompress(data, 16 + zlib.MAX_WBITS)
  386. elif encoding == "identity":
  387. return data
  388. else:
  389. raise RuntimeError(f"unknown content encoding: {encoding}")
  390. def _decode_content_transfer(self, data: bytes) -> bytes:
  391. encoding = self.headers.get(CONTENT_TRANSFER_ENCODING, "").lower()
  392. if encoding == "base64":
  393. return base64.b64decode(data)
  394. elif encoding == "quoted-printable":
  395. return binascii.a2b_qp(data)
  396. elif encoding in ("binary", "8bit", "7bit"):
  397. return data
  398. else:
  399. raise RuntimeError(
  400. "unknown content transfer encoding: {}" "".format(encoding)
  401. )
  402. def get_charset(self, default: str) -> str:
  403. """Returns charset parameter from Content-Type header or default."""
  404. ctype = self.headers.get(CONTENT_TYPE, "")
  405. mimetype = parse_mimetype(ctype)
  406. return mimetype.parameters.get("charset", default)
  407. @reify
  408. def name(self) -> Optional[str]:
  409. """Returns name specified in Content-Disposition header or None
  410. if missed or header is malformed.
  411. """
  412. _, params = parse_content_disposition(self.headers.get(CONTENT_DISPOSITION))
  413. return content_disposition_filename(params, "name")
  414. @reify
  415. def filename(self) -> Optional[str]:
  416. """Returns filename specified in Content-Disposition header or None
  417. if missed or header is malformed.
  418. """
  419. _, params = parse_content_disposition(self.headers.get(CONTENT_DISPOSITION))
  420. return content_disposition_filename(params, "filename")
  421. @payload_type(BodyPartReader, order=Order.try_first)
  422. class BodyPartReaderPayload(Payload):
  423. def __init__(self, value: BodyPartReader, *args: Any, **kwargs: Any) -> None:
  424. super().__init__(value, *args, **kwargs)
  425. params = {} # type: Dict[str, str]
  426. if value.name is not None:
  427. params["name"] = value.name
  428. if value.filename is not None:
  429. params["filename"] = value.filename
  430. if params:
  431. self.set_content_disposition("attachment", True, **params)
  432. async def write(self, writer: Any) -> None:
  433. field = self._value
  434. chunk = await field.read_chunk(size=2 ** 16)
  435. while chunk:
  436. await writer.write(field.decode(chunk))
  437. chunk = await field.read_chunk(size=2 ** 16)
  438. class MultipartReader:
  439. """Multipart body reader."""
  440. #: Response wrapper, used when multipart readers constructs from response.
  441. response_wrapper_cls = MultipartResponseWrapper
  442. #: Multipart reader class, used to handle multipart/* body parts.
  443. #: None points to type(self)
  444. multipart_reader_cls = None
  445. #: Body part reader class for non multipart/* content types.
  446. part_reader_cls = BodyPartReader
  447. def __init__(self, headers: Mapping[str, str], content: StreamReader) -> None:
  448. self.headers = headers
  449. self._boundary = ("--" + self._get_boundary()).encode()
  450. self._content = content
  451. self._last_part = (
  452. None
  453. ) # type: Optional[Union['MultipartReader', BodyPartReader]]
  454. self._at_eof = False
  455. self._at_bof = True
  456. self._unread = [] # type: List[bytes]
  457. def __aiter__(
  458. self,
  459. ) -> AsyncIterator["BodyPartReader"]:
  460. return self # type: ignore
  461. async def __anext__(
  462. self,
  463. ) -> Optional[Union["MultipartReader", BodyPartReader]]:
  464. part = await self.next()
  465. if part is None:
  466. raise StopAsyncIteration
  467. return part
  468. @classmethod
  469. def from_response(
  470. cls,
  471. response: "ClientResponse",
  472. ) -> MultipartResponseWrapper:
  473. """Constructs reader instance from HTTP response.
  474. :param response: :class:`~aiohttp.client.ClientResponse` instance
  475. """
  476. obj = cls.response_wrapper_cls(
  477. response, cls(response.headers, response.content)
  478. )
  479. return obj
  480. def at_eof(self) -> bool:
  481. """Returns True if the final boundary was reached or
  482. False otherwise.
  483. """
  484. return self._at_eof
  485. async def next(
  486. self,
  487. ) -> Optional[Union["MultipartReader", BodyPartReader]]:
  488. """Emits the next multipart body part."""
  489. # So, if we're at BOF, we need to skip till the boundary.
  490. if self._at_eof:
  491. return None
  492. await self._maybe_release_last_part()
  493. if self._at_bof:
  494. await self._read_until_first_boundary()
  495. self._at_bof = False
  496. else:
  497. await self._read_boundary()
  498. if self._at_eof: # we just read the last boundary, nothing to do there
  499. return None
  500. self._last_part = await self.fetch_next_part()
  501. return self._last_part
  502. async def release(self) -> None:
  503. """Reads all the body parts to the void till the final boundary."""
  504. while not self._at_eof:
  505. item = await self.next()
  506. if item is None:
  507. break
  508. await item.release()
  509. async def fetch_next_part(
  510. self,
  511. ) -> Union["MultipartReader", BodyPartReader]:
  512. """Returns the next body part reader."""
  513. headers = await self._read_headers()
  514. return self._get_part_reader(headers)
  515. def _get_part_reader(
  516. self,
  517. headers: "CIMultiDictProxy[str]",
  518. ) -> Union["MultipartReader", BodyPartReader]:
  519. """Dispatches the response by the `Content-Type` header, returning
  520. suitable reader instance.
  521. :param dict headers: Response headers
  522. """
  523. ctype = headers.get(CONTENT_TYPE, "")
  524. mimetype = parse_mimetype(ctype)
  525. if mimetype.type == "multipart":
  526. if self.multipart_reader_cls is None:
  527. return type(self)(headers, self._content)
  528. return self.multipart_reader_cls(headers, self._content)
  529. else:
  530. return self.part_reader_cls(self._boundary, headers, self._content)
  531. def _get_boundary(self) -> str:
  532. mimetype = parse_mimetype(self.headers[CONTENT_TYPE])
  533. assert mimetype.type == "multipart", "multipart/* content type expected"
  534. if "boundary" not in mimetype.parameters:
  535. raise ValueError(
  536. "boundary missed for Content-Type: %s" % self.headers[CONTENT_TYPE]
  537. )
  538. boundary = mimetype.parameters["boundary"]
  539. if len(boundary) > 70:
  540. raise ValueError("boundary %r is too long (70 chars max)" % boundary)
  541. return boundary
  542. async def _readline(self) -> bytes:
  543. if self._unread:
  544. return self._unread.pop()
  545. return await self._content.readline()
  546. async def _read_until_first_boundary(self) -> None:
  547. while True:
  548. chunk = await self._readline()
  549. if chunk == b"":
  550. raise ValueError(
  551. "Could not find starting boundary %r" % (self._boundary)
  552. )
  553. chunk = chunk.rstrip()
  554. if chunk == self._boundary:
  555. return
  556. elif chunk == self._boundary + b"--":
  557. self._at_eof = True
  558. return
  559. async def _read_boundary(self) -> None:
  560. chunk = (await self._readline()).rstrip()
  561. if chunk == self._boundary:
  562. pass
  563. elif chunk == self._boundary + b"--":
  564. self._at_eof = True
  565. epilogue = await self._readline()
  566. next_line = await self._readline()
  567. # the epilogue is expected and then either the end of input or the
  568. # parent multipart boundary, if the parent boundary is found then
  569. # it should be marked as unread and handed to the parent for
  570. # processing
  571. if next_line[:2] == b"--":
  572. self._unread.append(next_line)
  573. # otherwise the request is likely missing an epilogue and both
  574. # lines should be passed to the parent for processing
  575. # (this handles the old behavior gracefully)
  576. else:
  577. self._unread.extend([next_line, epilogue])
  578. else:
  579. raise ValueError(f"Invalid boundary {chunk!r}, expected {self._boundary!r}")
  580. async def _read_headers(self) -> "CIMultiDictProxy[str]":
  581. lines = [b""]
  582. while True:
  583. chunk = await self._content.readline()
  584. chunk = chunk.strip()
  585. lines.append(chunk)
  586. if not chunk:
  587. break
  588. parser = HeadersParser()
  589. headers, raw_headers = parser.parse_headers(lines)
  590. return headers
  591. async def _maybe_release_last_part(self) -> None:
  592. """Ensures that the last read body part is read completely."""
  593. if self._last_part is not None:
  594. if not self._last_part.at_eof():
  595. await self._last_part.release()
  596. self._unread.extend(self._last_part._unread)
  597. self._last_part = None
  598. _Part = Tuple[Payload, str, str]
  599. class MultipartWriter(Payload):
  600. """Multipart body writer."""
  601. def __init__(self, subtype: str = "mixed", boundary: Optional[str] = None) -> None:
  602. boundary = boundary if boundary is not None else uuid.uuid4().hex
  603. # The underlying Payload API demands a str (utf-8), not bytes,
  604. # so we need to ensure we don't lose anything during conversion.
  605. # As a result, require the boundary to be ASCII only.
  606. # In both situations.
  607. try:
  608. self._boundary = boundary.encode("ascii")
  609. except UnicodeEncodeError:
  610. raise ValueError("boundary should contain ASCII only chars") from None
  611. ctype = f"multipart/{subtype}; boundary={self._boundary_value}"
  612. super().__init__(None, content_type=ctype)
  613. self._parts = [] # type: List[_Part]
  614. def __enter__(self) -> "MultipartWriter":
  615. return self
  616. def __exit__(
  617. self,
  618. exc_type: Optional[Type[BaseException]],
  619. exc_val: Optional[BaseException],
  620. exc_tb: Optional[TracebackType],
  621. ) -> None:
  622. pass
  623. def __iter__(self) -> Iterator[_Part]:
  624. return iter(self._parts)
  625. def __len__(self) -> int:
  626. return len(self._parts)
  627. def __bool__(self) -> bool:
  628. return True
  629. _valid_tchar_regex = re.compile(br"\A[!#$%&'*+\-.^_`|~\w]+\Z")
  630. _invalid_qdtext_char_regex = re.compile(br"[\x00-\x08\x0A-\x1F\x7F]")
  631. @property
  632. def _boundary_value(self) -> str:
  633. """Wrap boundary parameter value in quotes, if necessary.
  634. Reads self.boundary and returns a unicode sting.
  635. """
  636. # Refer to RFCs 7231, 7230, 5234.
  637. #
  638. # parameter = token "=" ( token / quoted-string )
  639. # token = 1*tchar
  640. # quoted-string = DQUOTE *( qdtext / quoted-pair ) DQUOTE
  641. # qdtext = HTAB / SP / %x21 / %x23-5B / %x5D-7E / obs-text
  642. # obs-text = %x80-FF
  643. # quoted-pair = "\" ( HTAB / SP / VCHAR / obs-text )
  644. # tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*"
  645. # / "+" / "-" / "." / "^" / "_" / "`" / "|" / "~"
  646. # / DIGIT / ALPHA
  647. # ; any VCHAR, except delimiters
  648. # VCHAR = %x21-7E
  649. value = self._boundary
  650. if re.match(self._valid_tchar_regex, value):
  651. return value.decode("ascii") # cannot fail
  652. if re.search(self._invalid_qdtext_char_regex, value):
  653. raise ValueError("boundary value contains invalid characters")
  654. # escape %x5C and %x22
  655. quoted_value_content = value.replace(b"\\", b"\\\\")
  656. quoted_value_content = quoted_value_content.replace(b'"', b'\\"')
  657. return '"' + quoted_value_content.decode("ascii") + '"'
  658. @property
  659. def boundary(self) -> str:
  660. return self._boundary.decode("ascii")
  661. def append(self, obj: Any, headers: Optional[MultiMapping[str]] = None) -> Payload:
  662. if headers is None:
  663. headers = CIMultiDict()
  664. if isinstance(obj, Payload):
  665. obj.headers.update(headers)
  666. return self.append_payload(obj)
  667. else:
  668. try:
  669. payload = get_payload(obj, headers=headers)
  670. except LookupError:
  671. raise TypeError("Cannot create payload from %r" % obj)
  672. else:
  673. return self.append_payload(payload)
  674. def append_payload(self, payload: Payload) -> Payload:
  675. """Adds a new body part to multipart writer."""
  676. # compression
  677. encoding = payload.headers.get(
  678. CONTENT_ENCODING,
  679. "",
  680. ).lower() # type: Optional[str]
  681. if encoding and encoding not in ("deflate", "gzip", "identity"):
  682. raise RuntimeError(f"unknown content encoding: {encoding}")
  683. if encoding == "identity":
  684. encoding = None
  685. # te encoding
  686. te_encoding = payload.headers.get(
  687. CONTENT_TRANSFER_ENCODING,
  688. "",
  689. ).lower() # type: Optional[str]
  690. if te_encoding not in ("", "base64", "quoted-printable", "binary"):
  691. raise RuntimeError(
  692. "unknown content transfer encoding: {}" "".format(te_encoding)
  693. )
  694. if te_encoding == "binary":
  695. te_encoding = None
  696. # size
  697. size = payload.size
  698. if size is not None and not (encoding or te_encoding):
  699. payload.headers[CONTENT_LENGTH] = str(size)
  700. self._parts.append((payload, encoding, te_encoding)) # type: ignore
  701. return payload
  702. def append_json(
  703. self, obj: Any, headers: Optional[MultiMapping[str]] = None
  704. ) -> Payload:
  705. """Helper to append JSON part."""
  706. if headers is None:
  707. headers = CIMultiDict()
  708. return self.append_payload(JsonPayload(obj, headers=headers))
  709. def append_form(
  710. self,
  711. obj: Union[Sequence[Tuple[str, str]], Mapping[str, str]],
  712. headers: Optional[MultiMapping[str]] = None,
  713. ) -> Payload:
  714. """Helper to append form urlencoded part."""
  715. assert isinstance(obj, (Sequence, Mapping))
  716. if headers is None:
  717. headers = CIMultiDict()
  718. if isinstance(obj, Mapping):
  719. obj = list(obj.items())
  720. data = urlencode(obj, doseq=True)
  721. return self.append_payload(
  722. StringPayload(
  723. data, headers=headers, content_type="application/x-www-form-urlencoded"
  724. )
  725. )
  726. @property
  727. def size(self) -> Optional[int]:
  728. """Size of the payload."""
  729. total = 0
  730. for part, encoding, te_encoding in self._parts:
  731. if encoding or te_encoding or part.size is None:
  732. return None
  733. total += int(
  734. 2
  735. + len(self._boundary)
  736. + 2
  737. + part.size # b'--'+self._boundary+b'\r\n'
  738. + len(part._binary_headers)
  739. + 2 # b'\r\n'
  740. )
  741. total += 2 + len(self._boundary) + 4 # b'--'+self._boundary+b'--\r\n'
  742. return total
  743. async def write(self, writer: Any, close_boundary: bool = True) -> None:
  744. """Write body."""
  745. for part, encoding, te_encoding in self._parts:
  746. await writer.write(b"--" + self._boundary + b"\r\n")
  747. await writer.write(part._binary_headers)
  748. if encoding or te_encoding:
  749. w = MultipartPayloadWriter(writer)
  750. if encoding:
  751. w.enable_compression(encoding)
  752. if te_encoding:
  753. w.enable_encoding(te_encoding)
  754. await part.write(w) # type: ignore
  755. await w.write_eof()
  756. else:
  757. await part.write(writer)
  758. await writer.write(b"\r\n")
  759. if close_boundary:
  760. await writer.write(b"--" + self._boundary + b"--\r\n")
  761. class MultipartPayloadWriter:
  762. def __init__(self, writer: Any) -> None:
  763. self._writer = writer
  764. self._encoding = None # type: Optional[str]
  765. self._compress = None # type: Any
  766. self._encoding_buffer = None # type: Optional[bytearray]
  767. def enable_encoding(self, encoding: str) -> None:
  768. if encoding == "base64":
  769. self._encoding = encoding
  770. self._encoding_buffer = bytearray()
  771. elif encoding == "quoted-printable":
  772. self._encoding = "quoted-printable"
  773. def enable_compression(self, encoding: str = "deflate") -> None:
  774. zlib_mode = 16 + zlib.MAX_WBITS if encoding == "gzip" else -zlib.MAX_WBITS
  775. self._compress = zlib.compressobj(wbits=zlib_mode)
  776. async def write_eof(self) -> None:
  777. if self._compress is not None:
  778. chunk = self._compress.flush()
  779. if chunk:
  780. self._compress = None
  781. await self.write(chunk)
  782. if self._encoding == "base64":
  783. if self._encoding_buffer:
  784. await self._writer.write(base64.b64encode(self._encoding_buffer))
  785. async def write(self, chunk: bytes) -> None:
  786. if self._compress is not None:
  787. if chunk:
  788. chunk = self._compress.compress(chunk)
  789. if not chunk:
  790. return
  791. if self._encoding == "base64":
  792. buf = self._encoding_buffer
  793. assert buf is not None
  794. buf.extend(chunk)
  795. if buf:
  796. div, mod = divmod(len(buf), 3)
  797. enc_chunk, self._encoding_buffer = (buf[: div * 3], buf[div * 3 :])
  798. if enc_chunk:
  799. b64chunk = base64.b64encode(enc_chunk)
  800. await self._writer.write(b64chunk)
  801. elif self._encoding == "quoted-printable":
  802. await self._writer.write(binascii.b2a_qp(chunk))
  803. else:
  804. await self._writer.write(chunk)