import unittest

from cryptography.exceptions import InvalidSignature 
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
from cryptography.hazmat.primitives.kdf.hkdf import HKDF


class TestECC(unittest.TestCase):
	def test_dh(self):
		# slightly modified from:
		# https://cryptography.io/en/latest/hazmat/primitives/asymmetric/ec/#elliptic-curve-key-exchange-algorithm

		# private keys for use in the exchange.
		keya = b'0\x81\x87\x02\x01\x000\x13\x06\x07*\x86H\xce=\x02\x01\x06\x08*\x86H\xce=\x03\x01\x07\x04m0k\x02\x01\x01\x04 \x04\xa3X\xbd\x0e\xff*\x8cw\xf8\x9f\x05BD<\nY\xb3\xf1\xd2\xc1\xb0\r\x1e\xedu\x92]4M?\x01\xa1D\x03B\x00\x04P\xd9y\x92f\t\xa7x\xf3\xcf\x17O\xad\x93\xf9\x18"\t\xd3\x13*]3\xa7#\x8bH$j\xea\xfb\x8a\xd3\xb5\xee\xd9\x0f\x9c\xdb\xcc\xf1\xd7\x10\x88\x10e\x82-\x15CR\x08\xbe\x0c\x1e\x82p\x00C\xb2.O\x17\xd4'
		keyb = b'0\x81\x87\x02\x01\x000\x13\x06\x07*\x86H\xce=\x02\x01\x06\x08*\x86H\xce=\x03\x01\x07\x04m0k\x02\x01\x01\x04 \xfb\xf8\xf7\x9f\xa3\xb7\xed\x8cT@`\xf6\x9c\xbbv\x0e?\x87\xb1(\xf6\xa8\xb3`\x91\xb4\x92W\xc6\xaa\xf5~\xa1D\x03B\x00\x04\xb8\xfe\xe4\x8dkukc\xa4^\x87\x98\x9c\xb9\xa8\xec\x86\xf8\xc2\x89\xaeF\xe8q\xb9q\x92I\x98n\xfe\xe3<{\x1c&R\x82\xb1\x94=\xa5h*)m/\x13\xfb\x05\x1d\x98u\xec\x1ew\xdfW\x84\xfe\x9eSl\x83'

		#server_private_key = ec.generate_private_key(ec.SECP256R1())
		#print('spk:', repr(server_private_key.private_bytes(
		#    encoding=serialization.Encoding.DER,
		#    format=serialization.PublicFormat.UncompressedPoint,
		#    algorithm=None)))

		server_private_key = serialization.load_der_private_key(keya,
		    None, backend=default_backend())

		pubpoint = server_private_key.private_numbers().public_numbers

		self.assertEqual(server_private_key.private_numbers(). \
		    private_value,
		    2097859916579721232322403601989230314767884081400167668022347085958538411777)

		self.assertEqual(pubpoint.x,
		    36569272757924220784927781299997671003354453138026426189464987345196826688394)
		self.assertEqual(pubpoint.y,
		    95759458837377270694950453035845273914475242769728982836658929275588228093908)

		self.assertEqual(keya, server_private_key.private_bytes(
		    encoding=serialization.Encoding.DER,
		    format=serialization.PrivateFormat.PKCS8,
		    encryption_algorithm=serialization.NoEncryption()))

		# In a real handshake the peer is a remote client. For this
		# example we'll generate another local private key though.
		peer_private_key = serialization.load_der_private_key(keyb,
		    None)

		shared_key = server_private_key.exchange(
		    ec.ECDH(), peer_private_key.public_key())

		self.assertEqual(shared_key,
		    b'\x1a\x9f\x93c\xb0s\xa2\x15]{\xa3\xcc\xcf&Q\xd6g\x83\x86%\x7f\t\xfem@\xcb\xe9:U\x16\x07\x02')

		# Perform key derivation.
		derived_key = HKDF(algorithm=hashes.SHA256(),
		    length=32, salt=None, info=b'handshake data', backend=None
		    ).derive(shared_key)

		# And now we can demonstrate that the handshake performed in
		# the opposite direction gives the same final value

		same_shared_key = peer_private_key.exchange(
		    ec.ECDH(), server_private_key.public_key())

		self.assertEqual(shared_key, same_shared_key)

		# Perform key derivation.
		same_derived_key = HKDF(algorithm=hashes.SHA256(),
		    length=32, salt=None, info=b'handshake data',
		    ).derive(same_shared_key)

		self.assertEqual(derived_key,
		    b'\x89\r\xf7\xf0\xa6\xb9Z\xb9\xd7\xd0\x9b\x95y\xe0M\x11,\xb4\xe1Z\xe5\xa2j\xee)\xa0I\xb5Q\x18\x94\xd1')

		self.assertEqual(derived_key, same_derived_key)

	def test_decode_sig(self):
		sig = b"0D\x02 P\x92\xaf\xffoN\xadq\r=\x92\xb5\r\xe0l3\xf2\x80*\xdd|\xfe\xd8'\xb8\\\xe8\x94\xd6\xa1\xdb\xea\x02 \x18\x89j\xa8P\x83jk*\xb8\xa2\x15r&d\xa1\x9e\xf6\xec\xd2\xf4 \xd6\x08\x91bs\x18\xb5\x11/\x04"

		res = (36444202250238074078057463719437572015031876874679459625290239663827367091178, 11098302536735876471048588108325227764001515075886106999029234198831298588420)

		self.assertEqual(decode_dss_signature(sig), res)

	def test_sign(self):
		private_key = ec.generate_private_key(ec.SECP256R1())

		data = b"this is some data I'd like to sign"

		signature = private_key.sign(data, ec.ECDSA(hashes.SHA256()))

		# make sure we can decode our own signatures
		decode_dss_signature(signature)

		public_key = private_key.public_key()

		public_key.verify(signature, data, ec.ECDSA(hashes.SHA256()))

		wrongsig = bytearray(signature)
		wrongsig[0] ^= 1
		wrongsig[1] ^= 4
		wrongsig = bytes(wrongsig)

		self.assertRaises(InvalidSignature, public_key.verify, wrongsig, data, ec.ECDSA(hashes.SHA256()))

	def test_misc(self):
		keya = b'0\x81\x87\x02\x01\x000\x13\x06\x07*\x86H\xce=\x02\x01\x06\x08*\x86H\xce=\x03\x01\x07\x04m0k\x02\x01\x01\x04 \x04\xa3X\xbd\x0e\xff*\x8cw\xf8\x9f\x05BD<\nY\xb3\xf1\xd2\xc1\xb0\r\x1e\xedu\x92]4M?\x01\xa1D\x03B\x00\x04P\xd9y\x92f\t\xa7x\xf3\xcf\x17O\xad\x93\xf9\x18"\t\xd3\x13*]3\xa7#\x8bH$j\xea\xfb\x8a\xd3\xb5\xee\xd9\x0f\x9c\xdb\xcc\xf1\xd7\x10\x88\x10e\x82-\x15CR\x08\xbe\x0c\x1e\x82p\x00C\xb2.O\x17\xd4'
		skey = serialization.load_der_private_key(keya,
		    None, backend=default_backend())

		ckey = skey.public_key().public_bytes(
		    encoding=serialization.Encoding.X962,
		    format=serialization.PublicFormat.UncompressedPoint
		)

		self.assertEqual(ckey, b'\x04P\xd9y\x92f\t\xa7x\xf3\xcf\x17O\xad\x93\xf9\x18"\t\xd3\x13*]3\xa7#\x8bH$j\xea\xfb\x8a\xd3\xb5\xee\xd9\x0f\x9c\xdb\xcc\xf1\xd7\x10\x88\x10e\x82-\x15CR\x08\xbe\x0c\x1e\x82p\x00C\xb2.O\x17\xd4')

		pubkey = ec.EllipticCurvePublicKey.from_encoded_point(
		    ec.SECP256R1(), ckey)

		self.assertTrue(isinstance(pubkey, ec.EllipticCurvePublicKey))

		self.assertEqual(ckey, pubkey.public_bytes(
		    encoding=serialization.Encoding.X962,
		    format=serialization.PublicFormat.UncompressedPoint
		))