#!/usr/bin/env python
#
# Copyright 2015 clowwindy
#
# 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.

from __future__ import absolute_import, division, print_function, \
    with_statement

from ctypes import c_char_p, c_int, c_ulong, c_ulonglong, byref, \
    create_string_buffer, c_void_p

import sys

from shadowsocks.crypto import util

__all__ = ['ciphers']

libsodium = None
loaded = False

buf_size = 2048

# for salsa20 and chacha20 and chacha20-ietf
BLOCK_SIZE = 64


def load_libsodium():
    global loaded, libsodium, buf

    libsodium = util.find_library('sodium', 'crypto_stream_salsa20_xor_ic',
                                  'libsodium')
    if libsodium is None:
        raise Exception('libsodium not found')

    libsodium.crypto_stream_salsa20_xor_ic.restype = c_int
    libsodium.crypto_stream_salsa20_xor_ic.argtypes = (c_void_p, c_char_p,
                                                       c_ulonglong,
                                                       c_char_p, c_ulonglong,
                                                       c_char_p)
    libsodium.crypto_stream_chacha20_xor_ic.restype = c_int
    libsodium.crypto_stream_chacha20_xor_ic.argtypes = (c_void_p, c_char_p,
                                                        c_ulonglong,
                                                        c_char_p, c_ulonglong,
                                                        c_char_p)

    try:
        libsodium.crypto_stream_chacha20_ietf_xor_ic.restype = c_int
        libsodium.crypto_stream_chacha20_ietf_xor_ic.argtypes = (c_void_p, c_char_p,
                                                                 c_ulonglong,
                                                                 c_char_p, c_ulong,
                                                                 c_char_p)
    except:
        print("ChaCha20 IETF not support.", file=sys.stderr)
        pass

    try:
        libsodium.crypto_stream_xsalsa20_xor_ic.restype = c_int
        libsodium.crypto_stream_xsalsa20_xor_ic.argtypes = (c_void_p, c_char_p,
                                                            c_ulonglong,
                                                            c_char_p, c_ulonglong,
                                                            c_char_p)
    except:
        print("XSalsa20 not support.", file=sys.stderr)
        pass

    try:
        libsodium.crypto_stream_xchacha20_xor_ic.restype = c_int
        libsodium.crypto_stream_xchacha20_xor_ic.argtypes = (c_void_p, c_char_p,
                                                             c_ulonglong,
                                                             c_char_p, c_ulonglong,
                                                             c_char_p)
    except:
        print("XChaCha20 not support. XChaCha20 only support since libsodium v1.0.12", file=sys.stderr)
        pass

    buf = create_string_buffer(buf_size)
    loaded = True


class SodiumCrypto(object):
    def __init__(self, cipher_name, key, iv, op):
        if not loaded:
            load_libsodium()
        self.key = key
        self.iv = iv
        self.key_ptr = c_char_p(key)
        self.iv_ptr = c_char_p(iv)
        if cipher_name == 'salsa20':
            self.cipher = libsodium.crypto_stream_salsa20_xor_ic
        elif cipher_name == 'chacha20':
            self.cipher = libsodium.crypto_stream_chacha20_xor_ic
        elif cipher_name == 'chacha20-ietf':
            self.cipher = libsodium.crypto_stream_chacha20_ietf_xor_ic
        elif cipher_name == 'xchacha20':
            self.cipher = libsodium.crypto_stream_xchacha20_xor_ic
        elif cipher_name == 'xsalsa20':
            self.cipher = libsodium.crypto_stream_xsalsa20_xor_ic
        else:
            raise Exception('Unknown cipher')
        # byte counter, not block counter
        self.counter = 0

    def update(self, data):
        global buf_size, buf
        l = len(data)

        # we can only prepend some padding to make the encryption align to
        # blocks
        padding = self.counter % BLOCK_SIZE
        if buf_size < padding + l:
            buf_size = (padding + l) * 2
            buf = create_string_buffer(buf_size)

        if padding:
            data = (b'\0' * padding) + data
        self.cipher(byref(buf), c_char_p(data), padding + l,
                    self.iv_ptr, int(self.counter / BLOCK_SIZE), self.key_ptr)
        self.counter += l
        # buf is copied to a str object when we access buf.raw
        # strip off the padding
        return buf.raw[padding:padding + l]

    def clean(self):
        pass


ciphers = {
    'salsa20': (32, 8, SodiumCrypto),
    'chacha20': (32, 8, SodiumCrypto),
    'chacha20-ietf': (32, 12, SodiumCrypto),
    'xchacha20': (32, 24, SodiumCrypto),
    'xsalsa20': (32, 24, SodiumCrypto),
}


def test_salsa20():
    cipher = SodiumCrypto('salsa20', b'k' * 32, b'i' * 16, 1)
    decipher = SodiumCrypto('salsa20', b'k' * 32, b'i' * 16, 0)

    util.run_cipher(cipher, decipher)


def test_chacha20():
    cipher = SodiumCrypto('chacha20', b'k' * 32, b'i' * 16, 1)
    decipher = SodiumCrypto('chacha20', b'k' * 32, b'i' * 16, 0)

    util.run_cipher(cipher, decipher)


def test_chacha20_ietf():
    cipher = SodiumCrypto('chacha20-ietf', b'k' * 32, b'i' * 16, 1)
    decipher = SodiumCrypto('chacha20-ietf', b'k' * 32, b'i' * 16, 0)

    util.run_cipher(cipher, decipher)


def test_xchacha20():
    cipher = SodiumCrypto('xchacha20', b'k' * 32, b'i' * 24, 1)
    decipher = SodiumCrypto('xchacha20', b'k' * 32, b'i' * 24, 0)

    util.run_cipher(cipher, decipher)


def test_xsalsa20():
    cipher = SodiumCrypto('xsalsa20', b'k' * 32, b'i' * 24, 1)
    decipher = SodiumCrypto('xsalsa20', b'k' * 32, b'i' * 24, 0)

    util.run_cipher(cipher, decipher)


if __name__ == '__main__':
    test_chacha20_ietf()
    test_chacha20()
    test_salsa20()
    test_xchacha20()
    test_xsalsa20()