text: # Copyright The OpenTelemetry Authors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """
  AWS X-Ray Propagator
  --------------------
  
  The **AWS X-Ray Propagator** provides a propagator that when used, adds a `trace
  header`_ to outgoing traces that is compatible with the AWS X-Ray backend service.
  This allows the trace context to be propagated when a trace span multiple AWS
  services.
  
  Usage
  -----
  
  Use the provided AWS X-Ray Propagator to inject the necessary context into
  traces sent to external systems.
  
  This can be done by either setting this environment variable:
  
  ::
  
      export OTEL_PROPAGATORS = aws_xray
  
  
  Or by setting this propagator in your instrumented application:
  
  .. code-block:: python
  
      from opentelemetry.propagate import set_global_textmap
      from opentelemetry.sdk.extension.aws.trace.propagation.aws_xray_format import AwsXRayFormat
  
      set_global_textmap(AwsXRayFormat())
  
  API
  ---
  .. _trace header: https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-tracingheader
  """
  
  import logging
  import typing
  
  import opentelemetry.trace as trace
  from opentelemetry.context import Context
  from opentelemetry.propagators.textmap import (
      CarrierT,
      Getter,
      Setter,
      TextMapPropagator,
      default_getter,
      default_setter,
  )
  
  TRACE_HEADER_KEY = "X-Amzn-Trace-Id"
  KV_PAIR_DELIMITER = ";"
  KEY_AND_VALUE_DELIMITER = "="
  
  TRACE_ID_KEY = "Root"
  TRACE_ID_LENGTH = 35
  TRACE_ID_VERSION = "1"
  TRACE_ID_DELIMITER = "-"
  TRACE_ID_DELIMITER_INDEX_1 = 1
  TRACE_ID_DELIMITER_INDEX_2 = 10
  TRACE_ID_FIRST_PART_LENGTH = 8
  
  PARENT_ID_KEY = "Parent"
  PARENT_ID_LENGTH = 16
  
  SAMPLED_FLAG_KEY = "Sampled"
  SAMPLED_FLAG_LENGTH = 1
  IS_SAMPLED = "1"
  NOT_SAMPLED = "0"
  
  
  _logger = logging.getLogger(__name__)
  
  
  class AwsParseTraceHeaderError(Exception):
      def __init__(self, message):
          super().__init__()
          self.message = message
  
  
  class AwsXRayFormat(TextMapPropagator):
      """Propagator for the AWS X-Ray Trace Header propagation protocol.
  
      See: https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-tracingheader
      """
  
      # AWS
  
      def extract(
          self,
          carrier: CarrierT,
          context: typing.Optional[Context] = None,
          getter: Getter = default_getter,
      ) -> Context:
          trace_header_list = getter.get(carrier, TRACE_HEADER_KEY)
  
          if not trace_header_list or len(trace_header_list) != 1:
              return trace.set_span_in_context(
                  trace.INVALID_SPAN, context=context
              )
  
          trace_header = trace_header_list[0]
  
          if not trace_header:
              return trace.set_span_in_context(
                  trace.INVALID_SPAN, context=context
              )
  
          try:
              (
                  trace_id,
                  span_id,
                  sampled,
              ) = AwsXRayFormat._extract_span_properties(trace_header)
          except AwsParseTraceHeaderError as err:
              _logger.debug(err.message)
              return trace.set_span_in_context(
                  trace.INVALID_SPAN, context=context
              )
  
          options = 0
          if sampled:
              options |= trace.TraceFlags.SAMPLED
  
          span_context = trace.SpanContext(
              trace_id=trace_id,
              span_id=span_id,
              is_remote=True,
              trace_flags=trace.TraceFlags(options),
              trace_state=trace.TraceState(),
          )
  
          if not span_context.is_valid:
              _logger.debug(
                  "Invalid Span Extracted. Insertting INVALID span into provided context."
              )
              return trace.set_span_in_context(
                  trace.INVALID_SPAN, context=context
              )
  
          return trace.set_span_in_context(
              trace.NonRecordingSpan(span_context), context=context
          )
  
      @staticmethod
      def _extract_span_properties(trace_header):
          trace_id = trace.INVALID_TRACE_ID
          span_id = trace.INVALID_SPAN_ID
          sampled = False
  
          for kv_pair_str in trace_header.split(KV_PAIR_DELIMITER):
              try:
                  key_str, value_str = kv_pair_str.split(KEY_AND_VALUE_DELIMITER)
                  key, value = key_str.strip(), value_str.strip()
              except ValueError as ex:
                  raise AwsParseTraceHeaderError(
                      (
                          "Error parsing X-Ray trace header. Invalid key value pair: %s. Returning INVALID span context.",
                          kv_pair_str,
                      )
                  ) from ex
              if key == TRACE_ID_KEY:
                  if not AwsXRayFormat._validate_trace_id(value):
                      raise AwsParseTraceHeaderError(
                          (
                              "Invalid TraceId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
                              TRACE_HEADER_KEY,
                              trace_header,
                          )
                      )
  
                  try:
                      trace_id = AwsXRayFormat._parse_trace_id(value)
                  except ValueError as ex:
                      raise AwsParseTraceHeaderError(
                          (
                              "Invalid TraceId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
                              TRACE_HEADER_KEY,
                              trace_header,
                          )
                      ) from ex
              elif key == PARENT_ID_KEY:
                  if not AwsXRayFormat._validate_span_id(value):
                      raise AwsParseTraceHeaderError(
                          (
                              "Invalid ParentId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
                              TRACE_HEADER_KEY,
                              trace_header,
                          )
                      )
  
                  try:
                      span_id = AwsXRayFormat._parse_span_id(value)
                  except ValueError as ex:
                      raise AwsParseTraceHeaderError(
                          (
                              "Invalid TraceId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
                              TRACE_HEADER_KEY,
                              trace_header,
                          )
                      ) from ex
              elif key == SAMPLED_FLAG_KEY:
                  if not AwsXRayFormat._validate_sampled_flag(value):
                      raise AwsParseTraceHeaderError(
                          (
                              "Invalid Sampling flag in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
                              TRACE_HEADER_KEY,
                              trace_header,
                          )
                      )
  
                  sampled = AwsXRayFormat._parse_sampled_flag(value)
  
          return trace_id, span_id, sampled
  
      @staticmethod
      def _validate_trace_id(trace_id_str):
          return (
              len(trace_id_str) == TRACE_ID_LENGTH
              and trace_id_str.startswith(TRACE_ID_VERSION)
              and trace_id_str[TRACE_ID_DELIMITER_INDEX_1] == TRACE_ID_DELIMITER
              and trace_id_str[TRACE_ID_DELIMITER_INDEX_2] == TRACE_ID_DELIMITER
          )
  
      @staticmethod
      def _parse_trace_id(trace_id_str):
          timestamp_subset = trace_id_str[
              TRACE_ID_DELIMITER_INDEX_1 + 1 : TRACE_ID_DELIMITER_INDEX_2
          ]
          unique_id_subset = trace_id_str[
              TRACE_ID_DELIMITER_INDEX_2 + 1 : TRACE_ID_LENGTH
          ]
          return int(timestamp_subset + unique_id_subset, 16)
  
      @staticmethod
      def _validate_span_id(span_id_str):
          return len(span_id_str) == PARENT_ID_LENGTH
  
      @staticmethod
      def _parse_span_id(span_id_str):
          return int(span_id_str, 16)
  
      @staticmethod
      def _validate_sampled_flag(sampled_flag_str):
          return len(
              sampled_flag_str
          ) == SAMPLED_FLAG_LENGTH and sampled_flag_str in (
              IS_SAMPLED,
              NOT_SAMPLED,
          )
  
      @staticmethod
      def _parse_sampled_flag(sampled_flag_str):
          return sampled_flag_str[0] == IS_SAMPLED
  
      def inject(
          self,
          carrier: CarrierT,
          context: typing.Optional[Context] = None,
          setter: Setter = default_setter,
      ) -> None:
          span = trace.get_current_span(context=context)
  
          span_context = span.get_span_context()
          if not span_context.is_valid:
              return
  
          otel_trace_id = "{:032x}".format(span_context.trace_id)
          xray_trace_id = TRACE_ID_DELIMITER.join(
              [
                  TRACE_ID_VERSION,
                  otel_trace_id[:TRACE_ID_FIRST_PART_LENGTH],
                  otel_trace_id[TRACE_ID_FIRST_PART_LENGTH:],
              ]
          )
  
          parent_id = "{:016x}".format(span_context.span_id)
  
          sampling_flag = (
              IS_SAMPLED
              if span_context.trace_flags & trace.TraceFlags.SAMPLED
              else NOT_SAMPLED
          )
  
          # TODO: Add OT trace state to the X-Ray trace header
  
          trace_header = KV_PAIR_DELIMITER.join(
              [
                  KEY_AND_VALUE_DELIMITER.join([key, value])
                  for key, value in [
                      (TRACE_ID_KEY, xray_trace_id),
                      (PARENT_ID_KEY, parent_id),
                      (SAMPLED_FLAG_KEY, sampling_flag),
                  ]
              ]
          )
  
          setter.set(
              carrier, TRACE_HEADER_KEY, trace_header,
          )
  
      @property
      def fields(self):
          """Returns a set with the fields set in `inject`."""
  
          return {TRACE_HEADER_KEY}
  