#!/usr/bin/python3-dbg
# -*- coding: utf-8 -*-

import locale

from chrysacase import ChrysalideTestCase
from pychrysalide.analysis.contents import MemoryContent
from pychrysalide.analysis.scan import ContentScanner
from pychrysalide.analysis.scan import ScanOptions
from pychrysalide.analysis.scan.patterns.backends import AcismBackend
from pychrysalide import core
from pychrysalide.plugins.kaitai.parsers import KaitaiStruct
from pychrysalide.plugins.kaitai.rost import KaitaiTrigger


class TestScansWithKaitai(ChrysalideTestCase):
    """TestCase for ROST scan with the KaitaiStruct parsing."""

    @classmethod
    def setUpClass(cls):

        super(TestScansWithKaitai, cls).setUpClass()

        cls._options = ScanOptions()
        cls._options.backend_for_data = AcismBackend


    def testSimpleKaitaiDefinitionForScanning(self):
        """Rely on basic Kaitai simple definition for scanning."""

        definitions = '''
meta:
  id: basic_test
seq:
  - id: field0
    type: u1
'''

        kstruct = KaitaiStruct(definitions)

        trigger = KaitaiTrigger(kstruct)

        root_ns = core.get_rost_root_namespace()

        ns = root_ns.resolve('kaitai')
        ns.register_item(trigger)

        ns = ns.resolve('basic_test')
        self.assertEqual(ns, trigger)

        cnt = MemoryContent(b'\x01\x02\x03')

        rule = '''
rule testing {

   condition:
      kaitai.basic_test.field0 == 1

}
'''

        scanner = ContentScanner(rule)
        ctx = scanner.analyze(self._options, cnt)

        self.assertIsNotNone(ctx)

        self.assertFalse(ctx.has_match_for_rule('no_such_rule'))

        self.assertTrue(ctx.has_match_for_rule('testing'))


        definitions = '''
meta:
  id: other_basic_test
seq:
  - id: field0
    type: u1
  - id: field1
    type: u1
'''

        kstruct = KaitaiStruct(definitions)

        trigger = KaitaiTrigger(kstruct)

        root_ns = core.get_rost_root_namespace()

        ns = root_ns.resolve('kaitai')
        ns.register_item(trigger)

        ns = ns.resolve('other_basic_test')
        self.assertEqual(ns, trigger)

        cnt = MemoryContent(b'\x01\x02\x03')

        rule = '''
rule testing {

   condition:
      kaitai.other_basic_test.field0 == 1 and kaitai.other_basic_test.field1 == 2

}
'''

        scanner = ContentScanner(rule)
        ctx = scanner.analyze(self._options, cnt)

        self.assertIsNotNone(ctx)

        self.assertTrue(ctx.has_match_for_rule('testing'))


        rule = '''
rule testing {

   condition:
      kaitai.other_basic_test.field0 == 1 and kaitai.other_basic_testXXXX.field1 == 2

}
'''

        scanner = ContentScanner(rule)
        ctx = scanner.analyze(self._options, cnt)

        self.assertIsNotNone(ctx)

        self.assertFalse(ctx.has_match_for_rule('testing'))


    def testKaitaiDefinitionWithListForScanning(self):
        """Access list items from Kaitai definition when scanning with ROST."""

        definitions = '''
meta:
  id: test_with_list
seq:
  - id: field0
    type: u1
    repeat: eos
'''

        kstruct = KaitaiStruct(definitions)

        trigger = KaitaiTrigger(kstruct)

        root_ns = core.get_rost_root_namespace()

        ns = root_ns.resolve('kaitai')
        ns.register_item(trigger)

        ns = ns.resolve('test_with_list')
        self.assertEqual(ns, trigger)

        cnt = MemoryContent(b'\x01\x02\x03')

        rule = '''
rule testing {

   condition:
      kaitai.test_with_list.field0[0] == 1 and kaitai.test_with_list.field0[1] == 2 and kaitai.test_with_list.field0[2] == 3

}
'''

        scanner = ContentScanner(rule)
        ctx = scanner.analyze(self._options, cnt)

        self.assertIsNotNone(ctx)

        self.assertTrue(ctx.has_match_for_rule('testing'))