asn1.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419
  1. # This file is dual licensed under the terms of the Apache License, Version
  2. # 2.0, and the BSD License. See the LICENSE file in the root of this repository
  3. # for complete details.
  4. from __future__ import annotations
  5. import builtins
  6. import dataclasses
  7. import sys
  8. import types
  9. import typing
  10. if sys.version_info < (3, 11):
  11. import typing_extensions
  12. LiteralString = typing_extensions.LiteralString
  13. else:
  14. LiteralString = typing.LiteralString
  15. from cryptography.hazmat.bindings._rust import declarative_asn1
  16. if sys.version_info < (3, 10):
  17. NoneType = type(None)
  18. else:
  19. NoneType = types.NoneType # type: ignore[nonetype-type]
  20. T = typing.TypeVar("T", covariant=True)
  21. U = typing.TypeVar("U")
  22. Tag = typing.TypeVar("Tag", bound=LiteralString)
  23. @dataclasses.dataclass(frozen=True)
  24. class Variant(typing.Generic[U, Tag]):
  25. """
  26. A tagged variant for CHOICE fields with the same underlying type.
  27. Use this when you have multiple CHOICE alternatives with the same type
  28. and need to distinguish between them:
  29. foo: (
  30. Annotated[Variant[int, typing.Literal["IntA"]], Implicit(0)]
  31. | Annotated[Variant[int, typing.Literal["IntB"]], Implicit(1)]
  32. )
  33. Usage:
  34. example = Example(foo=Variant(5, "IntA"))
  35. decoded.foo.value # The int value
  36. decoded.foo.tag # "IntA" or "IntB"
  37. """
  38. value: U
  39. tag: str
  40. decode_der = declarative_asn1.decode_der
  41. encode_der = declarative_asn1.encode_der
  42. def _is_union(field_type: type) -> bool:
  43. # NOTE: types.UnionType for `T | U`, typing.Union for `Union[T, U]`.
  44. # TODO: Drop the `hasattr()` once the minimum supported Python version
  45. # is >= 3.10.
  46. union_types = (
  47. (types.UnionType, typing.Union)
  48. if hasattr(types, "UnionType")
  49. else (typing.Union,)
  50. )
  51. return typing.get_origin(field_type) in union_types
  52. def _extract_annotation(
  53. metadata: tuple, field_name: str
  54. ) -> declarative_asn1.Annotation:
  55. default = None
  56. encoding = None
  57. size = None
  58. for raw_annotation in metadata:
  59. if isinstance(raw_annotation, Default):
  60. if default is not None:
  61. raise TypeError(
  62. f"multiple DEFAULT annotations found in field "
  63. f"'{field_name}'"
  64. )
  65. default = raw_annotation.value
  66. elif isinstance(raw_annotation, declarative_asn1.Encoding):
  67. if encoding is not None:
  68. raise TypeError(
  69. f"multiple IMPLICIT/EXPLICIT annotations found in field "
  70. f"'{field_name}'"
  71. )
  72. encoding = raw_annotation
  73. elif isinstance(raw_annotation, declarative_asn1.Size):
  74. if size is not None:
  75. raise TypeError(
  76. f"multiple SIZE annotations found in field '{field_name}'"
  77. )
  78. size = raw_annotation
  79. else:
  80. raise TypeError(f"unsupported annotation: {raw_annotation}")
  81. return declarative_asn1.Annotation(
  82. default=default, encoding=encoding, size=size
  83. )
  84. def _normalize_field_type(
  85. field_type: typing.Any, field_name: str
  86. ) -> declarative_asn1.AnnotatedType:
  87. # Strip the `Annotated[...]` off, and populate the annotation
  88. # from it if it exists.
  89. if typing.get_origin(field_type) is typing.Annotated:
  90. annotation = _extract_annotation(field_type.__metadata__, field_name)
  91. field_type, *_ = typing.get_args(field_type)
  92. else:
  93. annotation = declarative_asn1.Annotation()
  94. if annotation.size is not None and (
  95. typing.get_origin(field_type) not in (builtins.list, SetOf)
  96. and field_type
  97. not in (
  98. builtins.bytes,
  99. builtins.str,
  100. BitString,
  101. IA5String,
  102. PrintableString,
  103. )
  104. ):
  105. raise TypeError(
  106. f"field '{field_name}' has a SIZE annotation, but SIZE "
  107. "annotations are only supported for fields of types: "
  108. "[SEQUENCE OF, SET OF, BIT STRING, OCTET STRING, UTF8String, "
  109. "PrintableString, IA5String]"
  110. )
  111. if field_type is TLV:
  112. if isinstance(annotation.encoding, Implicit):
  113. raise TypeError(
  114. f"field '{field_name}' has an IMPLICIT annotation, but "
  115. "IMPLICIT annotations are not supported for TLV types."
  116. )
  117. elif annotation.default is not None:
  118. raise TypeError(
  119. f"field '{field_name}' has a DEFAULT annotation, but "
  120. "DEFAULT annotations are not supported for TLV types."
  121. )
  122. if hasattr(field_type, "__asn1_root__"):
  123. root_type = field_type.__asn1_root__
  124. if not isinstance(
  125. root_type,
  126. (declarative_asn1.Type.Sequence, declarative_asn1.Type.Set),
  127. ):
  128. raise TypeError(f"unsupported root type: {root_type}")
  129. return declarative_asn1.AnnotatedType(
  130. typing.cast(declarative_asn1.Type, root_type), annotation
  131. )
  132. elif _is_union(field_type):
  133. union_args = typing.get_args(field_type)
  134. if len(union_args) == 2 and NoneType in union_args:
  135. # A Union between a type and None is an OPTIONAL
  136. optional_type = (
  137. union_args[0] if union_args[1] is type(None) else union_args[1]
  138. )
  139. if optional_type is TLV:
  140. raise TypeError(
  141. "optional TLV types (`TLV | None`) are not "
  142. "currently supported"
  143. )
  144. annotated_type = _normalize_field_type(optional_type, field_name)
  145. if not annotated_type.annotation.is_empty():
  146. raise TypeError(
  147. "optional (`X | None`) types cannot have `X` "
  148. "annotated: annotations must apply to the union "
  149. "(i.e: `Annotated[X | None, annotation]`)"
  150. )
  151. if annotation.default is not None:
  152. raise TypeError(
  153. "optional (`X | None`) types should not have a DEFAULT "
  154. "annotation"
  155. )
  156. rust_field_type = declarative_asn1.Type.Option(annotated_type)
  157. else:
  158. # Otherwise, the Union is a CHOICE
  159. if isinstance(annotation.encoding, Implicit):
  160. # CHOICEs cannot be IMPLICIT. See X.680 section 31.2.9.
  161. raise TypeError(
  162. "CHOICE (`X | Y | ...`) types should not have an IMPLICIT "
  163. "annotation"
  164. )
  165. variants = [
  166. _type_to_variant(arg, field_name)
  167. for arg in union_args
  168. if arg is not type(None)
  169. ]
  170. # Union types should either be all Variants
  171. # (`Variant[..] | Variant[..] | etc`) or all non Variants
  172. are_union_types_tagged = variants[0].tag_name is not None
  173. if any(
  174. (v.tag_name is not None) != are_union_types_tagged
  175. for v in variants
  176. ):
  177. raise TypeError(
  178. "When using `asn1.Variant` in a union, all the other "
  179. "types in the union must also be `asn1.Variant`"
  180. )
  181. if are_union_types_tagged:
  182. tags = {v.tag_name for v in variants}
  183. if len(variants) != len(tags):
  184. raise TypeError(
  185. "When using `asn1.Variant` in a union, the tags used "
  186. "must be unique"
  187. )
  188. rust_choice_type = declarative_asn1.Type.Choice(variants)
  189. # If None is part of the union types, this is an OPTIONAL CHOICE
  190. rust_field_type = (
  191. declarative_asn1.Type.Option(
  192. declarative_asn1.AnnotatedType(
  193. rust_choice_type, declarative_asn1.Annotation()
  194. )
  195. )
  196. if NoneType in union_args
  197. else rust_choice_type
  198. )
  199. elif typing.get_origin(field_type) is builtins.list:
  200. inner_type = _normalize_field_type(
  201. typing.get_args(field_type)[0], field_name
  202. )
  203. rust_field_type = declarative_asn1.Type.SequenceOf(inner_type)
  204. elif typing.get_origin(field_type) is SetOf:
  205. inner_type = _normalize_field_type(
  206. typing.get_args(field_type)[0], field_name
  207. )
  208. rust_field_type = declarative_asn1.Type.SetOf(inner_type)
  209. else:
  210. rust_field_type = declarative_asn1.non_root_python_to_rust(field_type)
  211. return declarative_asn1.AnnotatedType(rust_field_type, annotation)
  212. # Convert a type to a Variant. Used with types inside Union
  213. # annotations (T1, T2, etc in `Union[T1, T2, ...]`).
  214. def _type_to_variant(
  215. t: typing.Any, field_name: str
  216. ) -> declarative_asn1.Variant:
  217. is_annotated = typing.get_origin(t) is typing.Annotated
  218. inner_type = typing.get_args(t)[0] if is_annotated else t
  219. # Check if this is a Variant[T, Tag] type
  220. if typing.get_origin(inner_type) is Variant:
  221. value_type, tag_literal = typing.get_args(inner_type)
  222. if typing.get_origin(tag_literal) is not typing.Literal:
  223. raise TypeError(
  224. "When using `asn1.Variant` in a type annotation, the second "
  225. "type parameter must be a `typing.Literal` type. E.g: "
  226. '`Variant[int, typing.Literal["MyInt"]]`.'
  227. )
  228. tag_name = typing.get_args(tag_literal)[0]
  229. if hasattr(value_type, "__asn1_root__"):
  230. rust_type = value_type.__asn1_root__
  231. else:
  232. rust_type = declarative_asn1.non_root_python_to_rust(value_type)
  233. if is_annotated:
  234. ann_type = declarative_asn1.AnnotatedType(
  235. rust_type,
  236. _extract_annotation(t.__metadata__, field_name),
  237. )
  238. else:
  239. ann_type = declarative_asn1.AnnotatedType(
  240. rust_type,
  241. declarative_asn1.Annotation(),
  242. )
  243. return declarative_asn1.Variant(Variant, ann_type, tag_name)
  244. else:
  245. # Plain type (not a tagged Variant)
  246. return declarative_asn1.Variant(
  247. inner_type,
  248. _normalize_field_type(t, field_name),
  249. None,
  250. )
  251. def _annotate_fields(
  252. raw_fields: dict[str, type],
  253. ) -> dict[str, declarative_asn1.AnnotatedType]:
  254. fields = {}
  255. for field_name, field_type in raw_fields.items():
  256. # Recursively normalize the field type into something that the
  257. # Rust code can understand.
  258. annotated_field_type = _normalize_field_type(field_type, field_name)
  259. fields[field_name] = annotated_field_type
  260. return fields
  261. def _register_asn1_sequence(cls: type[U]) -> None:
  262. raw_fields = typing.get_type_hints(cls, include_extras=True)
  263. root = declarative_asn1.Type.Sequence(cls, _annotate_fields(raw_fields))
  264. setattr(cls, "__asn1_root__", root)
  265. def _register_asn1_set(cls: type[U]) -> None:
  266. raw_fields = typing.get_type_hints(cls, include_extras=True)
  267. root = declarative_asn1.Type.Set(cls, _annotate_fields(raw_fields))
  268. setattr(cls, "__asn1_root__", root)
  269. # Due to https://github.com/python/mypy/issues/19731, we can't define an alias
  270. # for `dataclass_transform` that conditionally points to `typing` or
  271. # `typing_extensions` depending on the Python version. We work around it by
  272. # making the whole decorated class conditional on the Python version.
  273. if sys.version_info < (3, 11):
  274. @typing_extensions.dataclass_transform(kw_only_default=True)
  275. def sequence(cls: type[U]) -> type[U]:
  276. # We use `dataclasses.dataclass` to add an __init__ method
  277. # to the class with keyword-only parameters.
  278. if sys.version_info >= (3, 10):
  279. dataclass_cls = dataclasses.dataclass(
  280. repr=False,
  281. eq=False,
  282. # `match_args` was added in Python 3.10 and defaults
  283. # to True
  284. match_args=False,
  285. # `kw_only` was added in Python 3.10 and defaults to
  286. # False
  287. kw_only=True,
  288. )(cls)
  289. else:
  290. dataclass_cls = dataclasses.dataclass(
  291. repr=False,
  292. eq=False,
  293. )(cls)
  294. _register_asn1_sequence(dataclass_cls)
  295. return dataclass_cls
  296. @typing_extensions.dataclass_transform(kw_only_default=True)
  297. def set(cls: type[U]) -> type[U]:
  298. # We use `dataclasses.dataclass` to add an __init__ method
  299. # to the class with keyword-only parameters.
  300. if sys.version_info >= (3, 10):
  301. dataclass_cls = dataclasses.dataclass(
  302. repr=False,
  303. eq=False,
  304. # `match_args` was added in Python 3.10 and defaults
  305. # to True
  306. match_args=False,
  307. # `kw_only` was added in Python 3.10 and defaults to
  308. # False
  309. kw_only=True,
  310. )(cls)
  311. else:
  312. dataclass_cls = dataclasses.dataclass(
  313. repr=False,
  314. eq=False,
  315. )(cls)
  316. _register_asn1_set(dataclass_cls)
  317. return dataclass_cls
  318. else:
  319. @typing.dataclass_transform(kw_only_default=True)
  320. def sequence(cls: type[U]) -> type[U]:
  321. # Only add an __init__ method, with keyword-only
  322. # parameters.
  323. dataclass_cls = dataclasses.dataclass(
  324. repr=False,
  325. eq=False,
  326. match_args=False,
  327. kw_only=True,
  328. )(cls)
  329. _register_asn1_sequence(dataclass_cls)
  330. return dataclass_cls
  331. @typing.dataclass_transform(kw_only_default=True)
  332. def set(cls: type[U]) -> type[U]:
  333. # Only add an __init__ method, with keyword-only
  334. # parameters.
  335. dataclass_cls = dataclasses.dataclass(
  336. repr=False,
  337. eq=False,
  338. match_args=False,
  339. kw_only=True,
  340. )(cls)
  341. _register_asn1_set(dataclass_cls)
  342. return dataclass_cls
  343. # TODO: replace with `Default[U]` once the min Python version is >= 3.12
  344. @dataclasses.dataclass(frozen=True)
  345. class Default(typing.Generic[U]):
  346. value: U
  347. SetOf = declarative_asn1.SetOf
  348. Explicit = declarative_asn1.Encoding.Explicit
  349. Implicit = declarative_asn1.Encoding.Implicit
  350. Size = declarative_asn1.Size
  351. PrintableString = declarative_asn1.PrintableString
  352. IA5String = declarative_asn1.IA5String
  353. UTCTime = declarative_asn1.UTCTime
  354. GeneralizedTime = declarative_asn1.GeneralizedTime
  355. BitString = declarative_asn1.BitString
  356. TLV = declarative_asn1.Tlv
  357. Null = declarative_asn1.Null