import nacl

from msg import RelayNetMsg

class RelayCell(RelayNetMsg):
    """All cells (which are sent inside a CircuitCellMsg, and so do not
    need their own circuitid) should be a subclass of this class."""


class StringCell(RelayCell):
    """Send an arbitrary string as a cell."""
    def __init__(self, str):
        self.data = str

    def __str__(self):
        return self.data.__str__()


class CloseCell(RelayCell):
    """Close the circuit this cell was sent on.  It should be sent
    _unencrypted_ (not within an EncryptedCell), and relays that receive
    one should forward it along the adjacent circuit, then close both
    the circuit it was received on and the adjacent one."""


class VanillaCreatedCircuitCell(RelayCell):
    """The message for responding to circuit creation in Vanilla Onion
    Routing."""

    def __init__(self, ntor_reply):
        self.ntor_reply = ntor_reply


class VanillaExtendCircuitCell(RelayCell):
    """The message for requesting circuit extension in Vanilla Onion
    Routing."""

    def __init__(self, hopaddr, ntor_request):
        self.hopaddr = hopaddr
        self.ntor_request = ntor_request


class VanillaExtendedCircuitCell(RelayCell):
    """The message for responding to circuit extension in Vanilla Onion
    Routing."""

    def __init__(self, ntor_reply):
        self.ntor_reply = ntor_reply


class TelescopingCreatedCircuitCell(RelayCell):
    """The message for responding to circuit creation in Telescoping Walking
    Onions."""

    def __init__(self, ntor_reply):
        self.ntor_reply = ntor_reply


class TelescopingExtendCircuitCell(RelayCell):
    """The message for requesting circuit extension in Telescoping Walking
    Onions."""

    def __init__(self, idx, ntor_request):
        self.idx = idx
        self.ntor_request = ntor_request


class TelescopingExtendedCircuitCell(RelayCell):
    """The message for responding to circuit extension in Telescoping Walking
    Onions."""

    def __init__(self, ntor_reply, snip):
        self.ntor_reply = ntor_reply
        self.snip = snip

class SinglePassCreatedCircuitCell(RelayCell):
    """The message for responding to circuit creation in Single-Pass Walking
    Onions."""

    def __init__(self, ntor_reply, encrypted_cell):
        self.ntor_reply = ntor_reply
        # The above field is sent in plaintext; the below is an
        # SinglePassCreatedEnc, or None if this is the last layer
        self.enc = encrypted_cell


class EncryptedCell(RelayCell):
    """Send a message encrypted with a symmetric key.  In this
    implementation, the encryption is not really done.  A hash of the
    key is stored with the message so that it can be checked at
    decryption time."""

    def __init__(self, key, msg):
        self.keyhash = nacl.hash.sha256(key)
        self.plaintext = msg

    def decrypt(self, key):
        keyhash = nacl.hash.sha256(key)
        if keyhash != self.keyhash:
            raise ValueError("EncryptedCell key mismatch")
        return self.plaintext

    def size(self):
        # Current Tor actually has no overhead for encryption
        return self.plaintext.size()


class SinglePassCreatedEnc(EncryptedCell):
    """The nested encrypted informaion inside a
    SinglePassCreatedCircuitCell.  nextlayer is itself a
    SinglePassCreatedCircuitCell."""
    def __init__(self, key, snip, vrf_output, nextlayer):
        super().__init__(key, (snip, vrf_output, nextlayer))
