asn1.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  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 dataclasses
  6. import sys
  7. import typing
  8. if sys.version_info < (3, 11):
  9. import typing_extensions
  10. # We use the `include_extras` parameter of `get_type_hints`, which was
  11. # added in Python 3.9. This can be replaced by the `typing` version
  12. # once the min version is >= 3.9
  13. if sys.version_info < (3, 9):
  14. get_type_hints = typing_extensions.get_type_hints
  15. else:
  16. get_type_hints = typing.get_type_hints
  17. else:
  18. get_type_hints = typing.get_type_hints
  19. from cryptography.hazmat.bindings._rust import declarative_asn1
  20. T = typing.TypeVar("T", covariant=True)
  21. U = typing.TypeVar("U")
  22. encode_der = declarative_asn1.encode_der
  23. def _normalize_field_type(
  24. field_type: typing.Any, field_name: str
  25. ) -> declarative_asn1.AnnotatedType:
  26. annotation = declarative_asn1.Annotation()
  27. if hasattr(field_type, "__asn1_root__"):
  28. annotated_root = field_type.__asn1_root__
  29. if not isinstance(annotated_root, declarative_asn1.AnnotatedType):
  30. raise TypeError(f"unsupported root type: {annotated_root}")
  31. return annotated_root
  32. else:
  33. rust_field_type = declarative_asn1.non_root_python_to_rust(field_type)
  34. return declarative_asn1.AnnotatedType(rust_field_type, annotation)
  35. def _annotate_fields(
  36. raw_fields: dict[str, type],
  37. ) -> dict[str, declarative_asn1.AnnotatedType]:
  38. fields = {}
  39. for field_name, field_type in raw_fields.items():
  40. # Recursively normalize the field type into something that the
  41. # Rust code can understand.
  42. annotated_field_type = _normalize_field_type(field_type, field_name)
  43. fields[field_name] = annotated_field_type
  44. return fields
  45. def _register_asn1_sequence(cls: type[U]) -> None:
  46. raw_fields = get_type_hints(cls, include_extras=True)
  47. root = declarative_asn1.AnnotatedType(
  48. declarative_asn1.Type.Sequence(cls, _annotate_fields(raw_fields)),
  49. declarative_asn1.Annotation(),
  50. )
  51. setattr(cls, "__asn1_root__", root)
  52. # Due to https://github.com/python/mypy/issues/19731, we can't define an alias
  53. # for `dataclass_transform` that conditionally points to `typing` or
  54. # `typing_extensions` depending on the Python version (like we do for
  55. # `get_type_hints`).
  56. # We work around it by making the whole decorated class conditional on the
  57. # Python version.
  58. if sys.version_info < (3, 11):
  59. @typing_extensions.dataclass_transform(kw_only_default=True)
  60. def sequence(cls: type[U]) -> type[U]:
  61. # We use `dataclasses.dataclass` to add an __init__ method
  62. # to the class with keyword-only parameters.
  63. if sys.version_info >= (3, 10):
  64. dataclass_cls = dataclasses.dataclass(
  65. repr=False,
  66. eq=False,
  67. # `match_args` was added in Python 3.10 and defaults
  68. # to True
  69. match_args=False,
  70. # `kw_only` was added in Python 3.10 and defaults to
  71. # False
  72. kw_only=True,
  73. )(cls)
  74. else:
  75. dataclass_cls = dataclasses.dataclass(
  76. repr=False,
  77. eq=False,
  78. )(cls)
  79. _register_asn1_sequence(dataclass_cls)
  80. return dataclass_cls
  81. else:
  82. @typing.dataclass_transform(kw_only_default=True)
  83. def sequence(cls: type[U]) -> type[U]:
  84. # Only add an __init__ method, with keyword-only
  85. # parameters.
  86. dataclass_cls = dataclasses.dataclass(
  87. repr=False,
  88. eq=False,
  89. match_args=False,
  90. kw_only=True,
  91. )(cls)
  92. _register_asn1_sequence(dataclass_cls)
  93. return dataclass_cls