#!/usr/bin/env python

import io

import os

import struct

from absl.testing import absltest

from cryptography import exceptions

from cryptography.hazmat.primitives.ciphers import aead

from google.protobuf import empty_pb2

from google.protobuf import timestamp_pb2

from google.protobuf import type_pb2

from google.protobuf import wrappers_pb2

from grr_api_client import utils

class MessageToFlatDictTest(absltest.TestCase):

def testEmpty(self):

dct = utils.MessageToFlatDict(empty_pb2.Empty(), lambda _, value: value)

self.assertEmpty(dct)

def testSingleField(self):

string = wrappers_pb2.StringValue()

string.value = "foo"

dct = utils.MessageToFlatDict(string, lambda _, value: value)

self.assertEqual(dct, {"value": "foo"})

def testMultipleFields(self):

timestamp = timestamp_pb2.Timestamp()

timestamp.seconds = 1337

timestamp.nanos = 42

dct = utils.MessageToFlatDict(timestamp, lambda _, value: value)

self.assertEqual(dct, {"seconds": 1337, "nanos": 42})

def testNestedMessages(self):

option = type_pb2.Option()

option.name = "foo"

option.value.type_url = "bar.baz.quux"

option.value.value = b"quux"

dct = utils.MessageToFlatDict(option, lambda _, value: value)

self.assertEqual(

dct,

{

"name": "foo",

"value.type_url": "bar.baz.quux",

"value.value": b"quux",

},

)

def testTransform(self):

i32 = wrappers_pb2.Int32Value()

i32.value = 1337

dct = utils.MessageToFlatDict(i32, lambda _, value: value * 2)

self.assertEqual(dct, {"value": 1337 * 2})

class GetCrowdstrikeDecodedBlobTest(absltest.TestCase):

def _CrowdstrikeEncode(self, data) -> bytes:

buf = io.BytesIO()

buf.write(b"CSQD")

buf.write(struct.pack("<Q", len(data)))

buf.write(utils.Xor(data, 0x7E))

return buf.getvalue()

def testDecodeSingleChunk(self):

content = b"foobarbaz"

iterator_content = iter((self._CrowdstrikeEncode(content),))

encoded_content = utils.BinaryChunkIterator(iterator_content)

decoded = encoded_content.DecodeCrowdStrikeQuarantineEncoding()

self.assertIsInstance(decoded, utils.BinaryChunkIterator)

self.assertEqual(b"".join(decoded), content)

def testDecode_RaisesIfCrowdstrikeIdentifierBytesMissing(self):

content = b"foobarbazbang"

iterator_content = iter((content,))

encoded_content = utils.BinaryChunkIterator(iterator_content)

with self.assertRaises(ValueError):

list(encoded_content.DecodeCrowdStrikeQuarantineEncoding())

def testDecodeSeveralChunks(self):

content = b"ABC" * 1024

encoded_content = self._CrowdstrikeEncode(content)

first_chunk = encoded_content[0:1024]

second_chunk = encoded_content[1024:2048]

third_chunk = encoded_content[2048:]

encoded_content = utils.BinaryChunkIterator(

iter((first_chunk, second_chunk, third_chunk))

)

decoded = encoded_content.DecodeCrowdStrikeQuarantineEncoding()

self.assertEqual(b"".join(decoded), content)

class AEADDecryptTest(absltest.TestCase):

def testReadExact(self):

key = os.urandom(32)

aesgcm = aead.AESGCM(key)

nonce = os.urandom(utils._AEAD_NONCE_SIZE)

adata = utils._AEAD_ADATA_FORMAT.pack(0, True)

encrypted = io.BytesIO(

nonce + aesgcm.encrypt(nonce, b"foobarbazquxnorf", adata)

)

decrypted = utils.AEADDecrypt(encrypted, key)

self.assertEqual(decrypted.read(3), b"foo")

self.assertEqual(decrypted.read(3), b"bar")

self.assertEqual(decrypted.read(3), b"baz")

self.assertEqual(decrypted.read(3), b"qux")

self.assertEqual(decrypted.read(4), b"norf")

self.assertEqual(decrypted.read(), b"")

def testIncorrectNonceLength(self):

key = os.urandom(32)

buf = io.BytesIO()

nonce = os.urandom(utils._AEAD_NONCE_SIZE - 1)

buf.write(nonce)

buf.seek(0, io.SEEK_SET)

with self.assertRaisesRegex(EOFError, "nonce length"):

utils.AEADDecrypt(buf, key).read()

def testIncorrectTag(self):

key = os.urandom(32)

aesgcm = aead.AESGCM(key)

buf = io.BytesIO()

nonce = os.urandom(utils._AEAD_NONCE_SIZE)

buf.write(nonce)

buf.write(aesgcm.encrypt(nonce, b"foo", b"QUUX"))

buf.seek(0, io.SEEK_SET)

with self.assertRaises(exceptions.InvalidTag):

utils.AEADDecrypt(buf, key).read()

def testIncorrectData(self):

key = os.urandom(32)

aesgcm = aead.AESGCM(key)

buf = io.BytesIO()

nonce = os.urandom(utils._AEAD_NONCE_SIZE)

adata = utils._AEAD_ADATA_FORMAT.pack(0, True)

buf.write(nonce)

buf.write(aesgcm.encrypt(nonce, b"foo", adata))

buf.getbuffer()[-1] ^= 0b10101010 # Corrupt last byte.

buf.seek(0, io.SEEK_SET)

with self.assertRaises(exceptions.InvalidTag):

utils.AEADDecrypt(buf, key).read()

if __name__ == "__main__":

absltest.main()