<commit_before>from __future__ import unicode_literals

from . import AuthenticatedMessage, NutsMessageTooLarge, NutsInvalidState
from .enums import ClientState, ServerState, Message
from .hkdf import HKDF
from .utils import ascii_bin, decode_version, encode_version, rng, encode_varint, decode_varint

from enum import Enum, IntEnum
from functools import partial
from itsdangerous import constant_time_compare
from logging import getLogger
from nacl.c import crypto_scalarmult
from nacl.public import PrivateKey
import cbor
import hashlib
import os
import sha3
import six
import socket
import string
import sys


_logger = getLogger('nuts.session')


def handshake_mac(*args):
    _logger.debug('Client handshake MACing %s (%d) with key %s (%d)' % (
        ascii_bin(b''.join(args[1:])),
        len(b''.join(args[1:])),
        ascii_bin(args[0]),
        len(args[0])))
    return hashlib.sha3_256(b''.join(args)).digest()[:8]


class Session(object):

    version = b'1.0'
    #: Maxiumum length of the sequence numbers. Used both for sanity checks of
    #: incoming data and calculatino of max MTU
    max_seqnum_length = 4

    def generate_and_set_session_key(self):
        self.session_key = HKDF(self.R_a + self.R_b, self.shared_key).expand(self.version, length=16)


    def handle(self, message):
        """ Message has been received from sender. It's assumed here that the link layer
        has filtered out messages not intended for us, or that has bit-errors.
        """
        if not message:
            return
        msg_type_byte = six.byte2int(message)
        valid_transitions = self.transition_map.get(self.state, [])
        if not msg_type_byte in valid_transitions:
            # Ignore
            _logger.info('Ignoring invalid transition')
            return

        handler = self.handlers.get(msg_type_byte)
        _logger.debug('Handling data %s' % ascii_bin(message))
        handler(message)


    def extract_message_data(self, data):
        """ Strips type, sequence numbers and MAC and returns the upper-layer protocol data. """
        seqnum_bytes = 0
        for byte in six.iterbytes(data[1:]):
            seqnum_bytes += 1
            if byte >> 7 == 0:
                break
        return data[1+seqnum_bytes:-self._mac_length]


    def validate_and_update_others_seqnum(self, data):
        seqnum_bytes = []
        for byte in six.iterbytes(data[1:]):
            seqnum_bytes.append(byte)
            if len(seqnum_bytes) > self.max_seqnum_length:
                return False
            if byte >> 7 == 0:
                break
        seqnum = decode_varint(seqnum_bytes)
        valid = seqnum >= self.other_seq
        if valid:
            self.other_seq = seqnum + 1
        return valid


    def init_session_mac(self, key, func_name, length):
        func = getattr(hashlib, func_name)
        self._mac_key = key
        self._mac_func = func
        self._mac_length = length
        _logger.info('Mac set to %s (%d)', func_name, length)


    def get_mac(self, *args, **kwargs):
        key = kwargs.get('key') or self._mac_key
        _logger.debug('Client MACing %s ', ascii_bin(b''.join(args)))
        return self._mac_func(key + b''.join(args)).digest()[:self._mac_length]


    def verify_mac(self, message):
        expected_mac = self.get_mac(message[:-self._mac_length])
        return constant_time_compare(message[-self._mac_length:], expected_mac)


    def send(self, data):
        """ Called from AuthChannel when it needs to send data through this
        session, or from the client-side application. This method adds type,
        sequence number and MAC.

        Throws a NutsInvalidState if the session is not established.
        """
        if not self.state == ClientState.established:
            raise NutsInvalidState("Session needs to be established to be able"
                " to send data, call .connect(address) first")
        if self.mtu and len(data) > self.mtu:
            raise NutsMessageTooLarge('Message of %d bytes is too long to be '
                'sent through this session, maximum size supported is %d '
                'bytes.' % (len(data), self.mtu))
        msg = six.int2byte(self.outgoing_command) + encode_varint(self.my_seq) + data
        self._send(msg + self.get_mac(msg))
        self.my_seq += 1


    @property
    def mtu(self):
        """ Get the maximum transmission unit (MTU) for this session. Returns
        None if there's no upper limit, or the maximum size in bytes.

        The MTU is calculated as the MTU of the channel minus the overhead of
        the session.
        """
        if self.channel.mtu:
            session_overhead = 1 + self._mac_length + self.max_seqnum_length
            session_mtu = self.channel.mtu - session_overhead
            return session_mtu
<commit_msg>Make sending data work on both client and server<commit_after>from __future__ import unicode_literals

from . import AuthenticatedMessage, NutsMessageTooLarge, NutsInvalidState
from .enums import ClientState, ServerState, Message
from .hkdf import HKDF
from .utils import ascii_bin, decode_version, encode_version, rng, encode_varint, decode_varint

from enum import Enum, IntEnum
from functools import partial
from itsdangerous import constant_time_compare
from logging import getLogger
from nacl.c import crypto_scalarmult
from nacl.public import PrivateKey
import cbor
import hashlib
import os
import sha3
import six
import socket
import string
import sys


_logger = getLogger('nuts.session')


def handshake_mac(*args):
    _logger.debug('Client handshake MACing %s (%d) with key %s (%d)' % (
        ascii_bin(b''.join(args[1:])),
        len(b''.join(args[1:])),
        ascii_bin(args[0]),
        len(args[0])))
    return hashlib.sha3_256(b''.join(args)).digest()[:8]


class Session(object):

    version = b'1.0'
    #: Maxiumum length of the sequence numbers. Used both for sanity checks of
    #: incoming data and calculatino of max MTU
    max_seqnum_length = 4

    def generate_and_set_session_key(self):
        self.session_key = HKDF(self.R_a + self.R_b, self.shared_key).expand(self.version, length=16)


    def handle(self, message):
        """ Message has been received from sender. It's assumed here that the link layer
        has filtered out messages not intended for us, or that has bit-errors.
        """
        if not message:
            return
        msg_type_byte = six.byte2int(message)
        valid_transitions = self.transition_map.get(self.state, [])
        if not msg_type_byte in valid_transitions:
            # Ignore
            _logger.info('Ignoring invalid transition')
            return

        handler = self.handlers.get(msg_type_byte)
        _logger.debug('Handling data %s' % ascii_bin(message))
        handler(message)


    def extract_message_data(self, data):
        """ Strips type, sequence numbers and MAC and returns the upper-layer protocol data. """
        seqnum_bytes = 0
        for byte in six.iterbytes(data[1:]):
            seqnum_bytes += 1
            if byte >> 7 == 0:
                break
        return data[1+seqnum_bytes:-self._mac_length]


    def validate_and_update_others_seqnum(self, data):
        seqnum_bytes = []
        for byte in six.iterbytes(data[1:]):
            seqnum_bytes.append(byte)
            if len(seqnum_bytes) > self.max_seqnum_length:
                return False
            if byte >> 7 == 0:
                break
        seqnum = decode_varint(seqnum_bytes)
        valid = seqnum >= self.other_seq
        if valid:
            self.other_seq = seqnum + 1
        return valid


    def init_session_mac(self, key, func_name, length):
        func = getattr(hashlib, func_name)
        self._mac_key = key
        self._mac_func = func
        self._mac_length = length
        _logger.info('Mac set to %s (%d)', func_name, length)


    def get_mac(self, *args, **kwargs):
        key = kwargs.get('key') or self._mac_key
        _logger.debug('Client MACing %s ', ascii_bin(b''.join(args)))
        return self._mac_func(key + b''.join(args)).digest()[:self._mac_length]


    def verify_mac(self, message):
        expected_mac = self.get_mac(message[:-self._mac_length])
        return constant_time_compare(message[-self._mac_length:], expected_mac)


    def send(self, data):
        """ Called from AuthChannel when it needs to send data through this
        session, or from the client-side application. This method adds type,
        sequence number and MAC.

        Throws a NutsInvalidState if the session is not established.
        """
        if not self.state in (ClientState.established, ServerState.established):
            raise NutsInvalidState("Session needs to be established to be able"
                " to send data, call .connect(address) first (was %s)" % self.state)
        if self.mtu and len(data) > self.mtu:
            raise NutsMessageTooLarge('Message of %d bytes is too long to be '
                'sent through this session, maximum size supported is %d '
                'bytes.' % (len(data), self.mtu))
        msg = six.int2byte(self.outgoing_command) + encode_varint(self.my_seq) + data
        self._send(msg + self.get_mac(msg))
        self.my_seq += 1


    @property
    def mtu(self):
        """ Get the maximum transmission unit (MTU) for this session. Returns
        None if there's no upper limit, or the maximum size in bytes.

        The MTU is calculated as the MTU of the channel minus the overhead of
        the session.
        """
        if self.channel.mtu:
            session_overhead = 1 + self._mac_length + self.max_seqnum_length
            session_mtu = self.channel.mtu - session_overhead
            return session_mtu

{'fraction_non_alphanumeric': 0.09151465464234475, 'fraction_numerical': 0.005943840951014552, 'mean_word_length': 2.6037666174298377, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7510098', 'n_tokens_mistral': 2952, 'n_tokens_neox': 2811, 'n_words': 855}