asn1.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533
  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 enum
  8. import sys
  9. import types
  10. import typing
  11. if sys.version_info < (3, 11):
  12. import typing_extensions
  13. LiteralString = typing_extensions.LiteralString
  14. else:
  15. LiteralString = typing.LiteralString
  16. from cryptography.hazmat.bindings._rust import declarative_asn1
  17. from cryptography.hazmat.bindings._rust import x509 as rust_x509
  18. if sys.version_info < (3, 10):
  19. NoneType = type(None)
  20. else:
  21. NoneType = types.NoneType # type: ignore[nonetype-type]
  22. T = typing.TypeVar("T", covariant=True)
  23. U = typing.TypeVar("U")
  24. Tag = typing.TypeVar("Tag", bound=LiteralString)
  25. @dataclasses.dataclass(frozen=True)
  26. class Variant(typing.Generic[U, Tag]):
  27. """
  28. A tagged variant for CHOICE fields with the same underlying type.
  29. Use this when you have multiple CHOICE alternatives with the same type
  30. and need to distinguish between them:
  31. foo: (
  32. Annotated[Variant[int, typing.Literal["IntA"]], Implicit(0)]
  33. | Annotated[Variant[int, typing.Literal["IntB"]], Implicit(1)]
  34. )
  35. Usage:
  36. example = Example(foo=Variant(5, "IntA"))
  37. decoded.foo.value # The int value
  38. decoded.foo.tag # "IntA" or "IntB"
  39. """
  40. value: U
  41. tag: str
  42. decode_der = declarative_asn1.decode_der
  43. encode_der = declarative_asn1.encode_der
  44. _X509_TYPES = (
  45. rust_x509.Certificate,
  46. rust_x509.CertificateSigningRequest,
  47. rust_x509.CertificateRevocationList,
  48. )
  49. def _check_x509_field_annotations(
  50. field_type: typing.Any,
  51. annotation: declarative_asn1.Annotation,
  52. field_name: str,
  53. ) -> None:
  54. if field_type in _X509_TYPES and isinstance(annotation.encoding, Implicit):
  55. raise TypeError(
  56. f"field '{field_name}' has an IMPLICIT annotation, but "
  57. "IMPLICIT annotations are not supported for X.509 types."
  58. )
  59. def _is_union(field_type: type) -> bool:
  60. # NOTE: types.UnionType for `T | U`, typing.Union for `Union[T, U]`.
  61. # TODO: Drop the `hasattr()` once the minimum supported Python version
  62. # is >= 3.10.
  63. union_types = (
  64. (types.UnionType, typing.Union)
  65. if hasattr(types, "UnionType")
  66. else (typing.Union,)
  67. )
  68. return typing.get_origin(field_type) in union_types
  69. def _resolve_type_aliases(field_type: typing.Any) -> typing.Any:
  70. # Recursively resolve PEP 695 (`type X = ...`) type aliases (Python
  71. # 3.12+) to their underlying value, so that the rest of the
  72. # normalization logic never encounters an alias. Aliases can refer
  73. # to other aliases and can appear at any level of nesting (e.g.
  74. # inside `Annotated[...]`, unions, or `list[...]`).
  75. if sys.version_info < (3, 12):
  76. return field_type
  77. while isinstance(field_type, typing.TypeAliasType):
  78. field_type = field_type.__value__
  79. args = typing.get_args(field_type)
  80. resolved_args = tuple(_resolve_type_aliases(arg) for arg in args)
  81. if resolved_args == args:
  82. # No aliases anywhere inside: return the type unchanged.
  83. return field_type
  84. if _is_union(field_type):
  85. # `X | Y` unions can't be rebuilt through their origin like
  86. # other generics below: `typing.get_origin` returns
  87. # `types.UnionType` for them, which is only subscriptable on
  88. # Python 3.14+. Rebuilding through `typing.Union` also
  89. # flattens any nested union introduced by an alias of a union
  90. # (e.g. `Time | int` where `type Time = UTCTime |
  91. # GeneralizedTime`), just like `typing.Union` would have done
  92. # if the alias had been written inline.
  93. return typing.Union[resolved_args]
  94. # An alias appeared inside a generic (e.g. `Annotated[Time, ...]`,
  95. # `list[MyInt]`, or `SetOf[MyInt]`): re-parameterize the generic
  96. # with the resolved arguments. Subscripting with a tuple is
  97. # equivalent to subscripting with multiple arguments.
  98. return typing.get_origin(field_type)[resolved_args]
  99. def _extract_annotation(
  100. metadata: tuple, field_name: str
  101. ) -> declarative_asn1.Annotation:
  102. default = None
  103. encoding = None
  104. size = None
  105. for raw_annotation in metadata:
  106. if isinstance(raw_annotation, Default):
  107. if default is not None:
  108. raise TypeError(
  109. f"multiple DEFAULT annotations found in field "
  110. f"'{field_name}'"
  111. )
  112. default = raw_annotation.value
  113. elif isinstance(raw_annotation, declarative_asn1.Encoding):
  114. if encoding is not None:
  115. raise TypeError(
  116. f"multiple IMPLICIT/EXPLICIT annotations found in field "
  117. f"'{field_name}'"
  118. )
  119. encoding = raw_annotation
  120. elif isinstance(raw_annotation, declarative_asn1.Size):
  121. if size is not None:
  122. raise TypeError(
  123. f"multiple SIZE annotations found in field '{field_name}'"
  124. )
  125. size = raw_annotation
  126. else:
  127. raise TypeError(f"unsupported annotation: {raw_annotation}")
  128. return declarative_asn1.Annotation(
  129. default=default, encoding=encoding, size=size
  130. )
  131. def _normalize_field_type(
  132. field_type: typing.Any, field_name: str
  133. ) -> declarative_asn1.AnnotatedType:
  134. field_type = _resolve_type_aliases(field_type)
  135. # Strip the `Annotated[...]` off, and populate the annotation
  136. # from it if it exists.
  137. if typing.get_origin(field_type) is typing.Annotated:
  138. annotation = _extract_annotation(field_type.__metadata__, field_name)
  139. field_type, *_ = typing.get_args(field_type)
  140. else:
  141. annotation = declarative_asn1.Annotation()
  142. if annotation.size is not None and (
  143. typing.get_origin(field_type) not in (builtins.list, SetOf)
  144. and field_type
  145. not in (
  146. builtins.bytes,
  147. builtins.str,
  148. BitString,
  149. IA5String,
  150. PrintableString,
  151. )
  152. ):
  153. raise TypeError(
  154. f"field '{field_name}' has a SIZE annotation, but SIZE "
  155. "annotations are only supported for fields of types: "
  156. "[SEQUENCE OF, SET OF, BIT STRING, OCTET STRING, UTF8String, "
  157. "PrintableString, IA5String]"
  158. )
  159. if field_type is TLV:
  160. if isinstance(annotation.encoding, Implicit):
  161. raise TypeError(
  162. f"field '{field_name}' has an IMPLICIT annotation, but "
  163. "IMPLICIT annotations are not supported for TLV types."
  164. )
  165. elif annotation.default is not None:
  166. raise TypeError(
  167. f"field '{field_name}' has a DEFAULT annotation, but "
  168. "DEFAULT annotations are not supported for TLV types."
  169. )
  170. _check_x509_field_annotations(field_type, annotation, field_name)
  171. if hasattr(field_type, "__asn1_root__"):
  172. root_type = field_type.__asn1_root__
  173. if not isinstance(
  174. root_type,
  175. (
  176. declarative_asn1.Type.Sequence,
  177. declarative_asn1.Type.Set,
  178. declarative_asn1.Type.ValueSet,
  179. ),
  180. ):
  181. raise TypeError(f"unsupported root type: {root_type}")
  182. return declarative_asn1.AnnotatedType(
  183. typing.cast(declarative_asn1.Type, root_type), annotation
  184. )
  185. elif _is_union(field_type):
  186. union_args = typing.get_args(field_type)
  187. if len(union_args) == 2 and NoneType in union_args:
  188. # A Union between a type and None is an OPTIONAL
  189. optional_type = (
  190. union_args[0] if union_args[1] is type(None) else union_args[1]
  191. )
  192. if optional_type is TLV:
  193. raise TypeError(
  194. "optional TLV types (`TLV | None`) are not "
  195. "currently supported"
  196. )
  197. # For optional types, the annotation is associated with the
  198. # union, so we check it against the inner type here.
  199. _check_x509_field_annotations(
  200. optional_type, annotation, field_name
  201. )
  202. annotated_type = _normalize_field_type(optional_type, field_name)
  203. if not annotated_type.annotation.is_empty():
  204. raise TypeError(
  205. "optional (`X | None`) types cannot have `X` "
  206. "annotated: annotations must apply to the union "
  207. "(i.e: `Annotated[X | None, annotation]`)"
  208. )
  209. if annotation.default is not None:
  210. raise TypeError(
  211. "optional (`X | None`) types should not have a DEFAULT "
  212. "annotation"
  213. )
  214. rust_field_type = declarative_asn1.Type.Option(annotated_type)
  215. else:
  216. # Otherwise, the Union is a CHOICE
  217. if isinstance(annotation.encoding, Implicit):
  218. # CHOICEs cannot be IMPLICIT. See X.680 section 31.2.9.
  219. raise TypeError(
  220. "CHOICE (`X | Y | ...`) types should not have an IMPLICIT "
  221. "annotation"
  222. )
  223. variants = [
  224. _type_to_variant(arg, field_name)
  225. for arg in union_args
  226. if arg is not type(None)
  227. ]
  228. # Union types should either be all Variants
  229. # (`Variant[..] | Variant[..] | etc`) or all non Variants
  230. are_union_types_tagged = variants[0].tag_name is not None
  231. if any(
  232. (v.tag_name is not None) != are_union_types_tagged
  233. for v in variants
  234. ):
  235. raise TypeError(
  236. "When using `asn1.Variant` in a union, all the other "
  237. "types in the union must also be `asn1.Variant`"
  238. )
  239. if are_union_types_tagged:
  240. tags = {v.tag_name for v in variants}
  241. if len(variants) != len(tags):
  242. raise TypeError(
  243. "When using `asn1.Variant` in a union, the tags used "
  244. "must be unique"
  245. )
  246. rust_choice_type = declarative_asn1.Type.Choice(variants)
  247. # If None is part of the union types, this is an OPTIONAL CHOICE
  248. rust_field_type = (
  249. declarative_asn1.Type.Option(
  250. declarative_asn1.AnnotatedType(
  251. rust_choice_type, declarative_asn1.Annotation()
  252. )
  253. )
  254. if NoneType in union_args
  255. else rust_choice_type
  256. )
  257. elif typing.get_origin(field_type) is builtins.list:
  258. inner_type = _normalize_field_type(
  259. typing.get_args(field_type)[0], field_name
  260. )
  261. rust_field_type = declarative_asn1.Type.SequenceOf(inner_type)
  262. elif typing.get_origin(field_type) is SetOf:
  263. inner_type = _normalize_field_type(
  264. typing.get_args(field_type)[0], field_name
  265. )
  266. rust_field_type = declarative_asn1.Type.SetOf(inner_type)
  267. else:
  268. rust_field_type = declarative_asn1.non_root_python_to_rust(field_type)
  269. return declarative_asn1.AnnotatedType(rust_field_type, annotation)
  270. # Convert a type to a Variant. Used with types inside Union
  271. # annotations (T1, T2, etc in `Union[T1, T2, ...]`).
  272. def _type_to_variant(
  273. t: typing.Any, field_name: str
  274. ) -> declarative_asn1.Variant:
  275. is_annotated = typing.get_origin(t) is typing.Annotated
  276. inner_type = typing.get_args(t)[0] if is_annotated else t
  277. # Check if this is a Variant[T, Tag] type
  278. if typing.get_origin(inner_type) is Variant:
  279. value_type, tag_literal = typing.get_args(inner_type)
  280. if typing.get_origin(tag_literal) is not typing.Literal:
  281. raise TypeError(
  282. "When using `asn1.Variant` in a type annotation, the second "
  283. "type parameter must be a `typing.Literal` type. E.g: "
  284. '`Variant[int, typing.Literal["MyInt"]]`.'
  285. )
  286. tag_name = typing.get_args(tag_literal)[0]
  287. if hasattr(value_type, "__asn1_root__"):
  288. rust_type = value_type.__asn1_root__
  289. else:
  290. rust_type = declarative_asn1.non_root_python_to_rust(value_type)
  291. if is_annotated:
  292. ann_type = declarative_asn1.AnnotatedType(
  293. rust_type,
  294. _extract_annotation(t.__metadata__, field_name),
  295. )
  296. else:
  297. ann_type = declarative_asn1.AnnotatedType(
  298. rust_type,
  299. declarative_asn1.Annotation(),
  300. )
  301. return declarative_asn1.Variant(Variant, ann_type, tag_name)
  302. else:
  303. # Plain type (not a tagged Variant)
  304. return declarative_asn1.Variant(
  305. inner_type,
  306. _normalize_field_type(t, field_name),
  307. None,
  308. )
  309. def _annotate_fields(
  310. raw_fields: dict[str, type],
  311. ) -> dict[str, declarative_asn1.AnnotatedType]:
  312. fields = {}
  313. for field_name, field_type in raw_fields.items():
  314. # Recursively normalize the field type into something that the
  315. # Rust code can understand.
  316. annotated_field_type = _normalize_field_type(field_type, field_name)
  317. fields[field_name] = annotated_field_type
  318. return fields
  319. def _register_asn1_sequence(cls: type[U]) -> None:
  320. raw_fields = typing.get_type_hints(cls, include_extras=True)
  321. root = declarative_asn1.Type.Sequence(cls, _annotate_fields(raw_fields))
  322. setattr(cls, "__asn1_root__", root)
  323. def _register_asn1_set(cls: type[U]) -> None:
  324. raw_fields = typing.get_type_hints(cls, include_extras=True)
  325. root = declarative_asn1.Type.Set(cls, _annotate_fields(raw_fields))
  326. setattr(cls, "__asn1_root__", root)
  327. # Due to https://github.com/python/mypy/issues/19731, we can't define an alias
  328. # for `dataclass_transform` that conditionally points to `typing` or
  329. # `typing_extensions` depending on the Python version. We work around it by
  330. # making the whole decorated class conditional on the Python version.
  331. if sys.version_info < (3, 11):
  332. @typing_extensions.dataclass_transform(kw_only_default=True)
  333. def sequence(cls: type[U]) -> type[U]:
  334. # We use `dataclasses.dataclass` to add an __init__ method
  335. # to the class with keyword-only parameters.
  336. if sys.version_info >= (3, 10):
  337. dataclass_cls = dataclasses.dataclass(
  338. repr=False,
  339. eq=False,
  340. # `match_args` was added in Python 3.10 and defaults
  341. # to True
  342. match_args=False,
  343. # `kw_only` was added in Python 3.10 and defaults to
  344. # False
  345. kw_only=True,
  346. )(cls)
  347. else:
  348. dataclass_cls = dataclasses.dataclass(
  349. repr=False,
  350. eq=False,
  351. )(cls)
  352. _register_asn1_sequence(dataclass_cls)
  353. return dataclass_cls
  354. @typing_extensions.dataclass_transform(kw_only_default=True)
  355. def set(cls: type[U]) -> type[U]:
  356. # We use `dataclasses.dataclass` to add an __init__ method
  357. # to the class with keyword-only parameters.
  358. if sys.version_info >= (3, 10):
  359. dataclass_cls = dataclasses.dataclass(
  360. repr=False,
  361. eq=False,
  362. # `match_args` was added in Python 3.10 and defaults
  363. # to True
  364. match_args=False,
  365. # `kw_only` was added in Python 3.10 and defaults to
  366. # False
  367. kw_only=True,
  368. )(cls)
  369. else:
  370. dataclass_cls = dataclasses.dataclass(
  371. repr=False,
  372. eq=False,
  373. )(cls)
  374. _register_asn1_set(dataclass_cls)
  375. return dataclass_cls
  376. else:
  377. @typing.dataclass_transform(kw_only_default=True)
  378. def sequence(cls: type[U]) -> type[U]:
  379. # Only add an __init__ method, with keyword-only
  380. # parameters.
  381. dataclass_cls = dataclasses.dataclass(
  382. repr=False,
  383. eq=False,
  384. match_args=False,
  385. kw_only=True,
  386. )(cls)
  387. _register_asn1_sequence(dataclass_cls)
  388. return dataclass_cls
  389. @typing.dataclass_transform(kw_only_default=True)
  390. def set(cls: type[U]) -> type[U]:
  391. # Only add an __init__ method, with keyword-only
  392. # parameters.
  393. dataclass_cls = dataclasses.dataclass(
  394. repr=False,
  395. eq=False,
  396. match_args=False,
  397. kw_only=True,
  398. )(cls)
  399. _register_asn1_set(dataclass_cls)
  400. return dataclass_cls
  401. def value_set(
  402. value_type: type,
  403. ) -> typing.Callable[[type[U]], type[U]]:
  404. """
  405. A class decorator that registers an `enum.Enum` subclass as an
  406. ASN.1 value set of the given underlying type. All the member
  407. values must be instances of `value_type`. Members are encoded as
  408. their value; decoding fails if the decoded value does not match
  409. any member.
  410. """
  411. rust_type = declarative_asn1.non_root_python_to_rust(value_type)
  412. def decorator(cls: type[U]) -> type[U]:
  413. if not issubclass(cls, enum.Enum):
  414. raise TypeError(
  415. "value sets can only be defined from enum.Enum subclasses"
  416. )
  417. members = list(cls)
  418. if not members:
  419. raise TypeError(
  420. f"value set '{cls.__name__}' must have at least one member"
  421. )
  422. for member in members:
  423. if not isinstance(member.value, value_type):
  424. raise TypeError(
  425. f"member '{member.name}' of value set '{cls.__name__}' "
  426. f"must have a value of type "
  427. f"'{value_type.__name__}', got: "
  428. f"'{type(member.value).__name__}'"
  429. )
  430. inner = declarative_asn1.AnnotatedType(
  431. rust_type, declarative_asn1.Annotation()
  432. )
  433. # Map from member value to member, used for O(1) lookups when
  434. # decoding. This requires the member values to be hashable.
  435. value_map = {member.value: member for member in members}
  436. root = declarative_asn1.Type.ValueSet(cls, inner, value_map)
  437. setattr(cls, "__asn1_root__", root)
  438. return cls
  439. return decorator
  440. # TODO: replace with `Default[U]` once the min Python version is >= 3.12
  441. @dataclasses.dataclass(frozen=True)
  442. class Default(typing.Generic[U]):
  443. value: U
  444. SetOf = declarative_asn1.SetOf
  445. Explicit = declarative_asn1.Encoding.Explicit
  446. Implicit = declarative_asn1.Encoding.Implicit
  447. Size = declarative_asn1.Size
  448. PrintableString = declarative_asn1.PrintableString
  449. IA5String = declarative_asn1.IA5String
  450. UTCTime = declarative_asn1.UTCTime
  451. GeneralizedTime = declarative_asn1.GeneralizedTime
  452. BitString = declarative_asn1.BitString
  453. TLV = declarative_asn1.Tlv
  454. Null = declarative_asn1.Null