# Source code for pathspider.chains.dscp

    """
    .. module:: pathspider.chains.dscp
       :synopsis: A flow analysis chain for Differentiated Services

    This module contains the DSCPChain flow analysis chain which can be used by
    PATHspider's Observer for recording Differentiated Services [RFC2474]_ details.

    .. codeauthor:: Iain R. Learmonth <irl@fsfe.org>

    """

    from pathspider.chains.base import Chain
    from pathspider.chains.tcp import TCP_SYN

    class DSCPChain(Chain):
        """
        This flow analysis chain records details of the Differentiated Services
        Field in the IP header.

        +-----------------------+------+------------------------------------------------+
        | Field Name            | Type | Meaning                                        |
        +=======================+======+================================================+
        | ``dscp_mark_syn_fwd`` | int  | The value of the Differentiated Services       |
        |                       |      | codepoint seen on a TCP SYN packet in the      |
        |                       |      | forward direction                              |
        +-----------------------+------+------------------------------------------------+
        | ``dscp_mark_data_fwd``| int  | The value of the Differentiated Services       |
        |                       |      | codepoint seen on a non-TCP packet or a TCP    |
        |                       |      | packet with a payload in the forward direction |
        +-----------------------+------+------------------------------------------------+
        | ``dscp_mark_syn_rev`` | int  | The value of the Differentiated Services       |
        |                       |      | codepoint seen on a TCP SYN packet in the      |
        |                       |      | reverse direction                              |
        +-----------------------+------+------------------------------------------------+
        | ``dscp_mark_data_rev``| int  | The value of the Differentiated Services       |
        |                       |      | codepoint seen on a non-TCP packet or a TCP    |
        |                       |      | packet with a payload in the reverse direction |
        +-----------------------+------+------------------------------------------------+
        """

        def new_flow(self, rec, ip):
            """
            For a new flow, all fields will be initialised to ``None``.

            :param rec: the flow record
            :type rec: dict
            :param ip: the IP or IPv6 packet that triggered the creation of a new
                       flow record
            :type ip: plt.ip or plt.ip6
            :return: Always ``True``
            :rtype: bool
            """

            rec['dscp_mark_syn_fwd'] = None
            rec['dscp_mark_syn_rev'] = None
            rec['dscp_mark_data_fwd'] = None
            rec['dscp_mark_data_rev'] = None
            return True

        def ip4(self, rec, ip, rev):
            """
            Records DSCP markings from an IPv4 header.

            DSCP Marking
                For the first TCP SYN packet and the first non-TCP packet or TCP
                packet with a payload, the DSCP value will be recorded in the
                relevant field.

            :param rec: the flow record
            :type rec: dict
            :param ip: the IPv4 packet that was observed to be part of this flow
            :type ip: plt.ip
            :param rev: True if the packet was in the reverse direction, False if
                        in the forward direction
            :type rev: bool
            :return: Always ``True``
            :rtype: bool
            """

            return self._dscp_extract(rec, ip, rev)

        def ip6(self, rec, ip, rev):
            """
            Records DSCP markings from an IPv6 header.

            DSCP Marking
                For the first TCP SYN packet and the first non-TCP packet or TCP
                packet with a payload, the DSCP value will be recorded in the
                relevant field.

            :param rec: the flow record
            :type rec: dict
            :param ip: the IPv6 packet that was observed to be part of this flow
            :type ip: plt.ip6
            :param rev: True if the packet was in the reverse direction, False if
                        in the forward direction
            :type rev: bool
            :return: Always ``True``
            :rtype: bool
            """

            return self._dscp_extract(rec, ip, rev)

        def _dscp_extract(self, rec, ip, rev):
            tos = ip.traffic_class
            dscp = tos >> 2

            if ip.tcp:
                if ip.tcp.flags & TCP_SYN == TCP_SYN:
                    rec['dscp_mark_syn_rev' if rev else 'dscp_mark_syn_fwd'] = dscp
                    return True
                if ip.tcp.payload is None:
                    return True

            # If not TCP or TCP non-SYN
            data_key = 'dscp_mark_data_rev' if rev else 'dscp_mark_data_fwd'
            rec[data_key] = rec[data_key] or dscp
            return True