# -*- coding: utf-8 -*-

"""
    Aho-Corasick string search algorithm.

    Author    : Wojciech Muła, wojciech_mula@poczta.onet.pl
    WWW       : http://0x80.pl
    License   : public domain
"""

import ahocorasick
import unittest
import struct

from unittest.case import skipIf


class TreeNodeBuilderBase(object):

    def __init__(self):
        self.integer = 0
        self.fail = 0
        self.n = 0
        self.eow = 0
        self.next = []

    def dump(self):

        assert self.n == len(self.next)

        next = b''
        for letter, node in self.next:
            next += self.dump_edge(letter, node)

        return self.dump_node() + next


class TreeNodeBuilderPy3(TreeNodeBuilderBase):

    def dump_node(self):
        """
        On Debian 64-bit, GCC 7.3

        python3:

            integer   : size 8, offset 0
            fail      : size 8, offset 8
            n         : size 4, offset 16
            eow       : size 1, offset 20
            padding   : size 3
            next      : size 8, offset 24 -- omitted in dump
        """
        node = struct.pack('=QQIBxxx',
            self.integer,
            self.fail,
            self.n,
            self.eow)

        assert len(node) == 24

        return node

    def dump_edge(self, letter, node):
        assert ord(letter) < 256

        b = struct.pack('=IQ', ord(letter), node)
        assert len(b) == 12

        return b


TreeNodeBuilder = TreeNodeBuilderPy3

USE_EXACT_RAW = True


class TestUnpickleRaw(unittest.TestCase):

    def __init__(self, *args):
        super(TestUnpickleRaw, self).__init__(*args)

    # raw constructor get 7-tuple (see Automaton.c):
    # 1. serialized nodes (as list of bytes or strings)
    # 2. kind
    # 3. store
    # 4. key type
    # 5. word count
    # 6. length of the longest word
    # 7. python values saved in a trie (if store == ahocorasick.STORE_ANY)

    def setUp(self):
        self.count = 0
        self.raw = b''
        self.kind = ahocorasick.EMPTY
        self.store = ahocorasick.STORE_ANY
        self.key_type = ahocorasick.KEY_STRING
        self.word_count = 0
        self.longest = 0
        self.values = []

    def create_automaton(self, use_exact_raw=False):
        # alter values that were set in setUp
        if use_exact_raw:
            raw = self.raw
        else:
            raw = [self.create_raw_count(self.count) + self.raw]

        args = (raw, self.kind, self.store, self.key_type,
                self.word_count, self.longest, self.values);

        return ahocorasick.Automaton(*args)

    def create_node_builder(self, eow, children):
        builder = TreeNodeBuilder()
        builder.next = [(letter, i + 1) for letter, i in children]  # starts from 1
        builder.n = len(children)
        builder.eow = eow

        return builder

    def create_raw_count(self, n):
        return struct.pack('Q', n)

    def create_raw_node(self, eow, children):
        return self.create_node_builder(eow, children).dump()

    # --------------------------------------------------

    def test__construct_empty(self):

        A = self.create_automaton()

        self.assertTrue(A.kind == ahocorasick.EMPTY)
        self.assertTrue(len(A) == 0)

    @skipIf(not ahocorasick.unicode, "Run only with unicode build")
    def test__construct_simple_trie(self):

        r"""
        trie for set {he, her, his, him, it}

        #0 -> [h #1 ] -> [e #2*] -> [r #3*]
         |           \-> [i #4 ] -> [s #5*]
         |                      \-> [m #6*]
         |
         +--> [i #7 ] -> [t #8 ]
        """
        values = ["HE", "HER", "HIS", "HIM", "IT"]

        node0 = self.create_raw_node(0, [('h', 1), ('i', 7)])
        node1 = self.create_raw_node(0, [('e', 2), ('i', 4)])
        node2 = self.create_raw_node(1, [('r', 3)])  # HE
        node3 = self.create_raw_node(1, [])  # HER
        node4 = self.create_raw_node(0, [('s', 5), ('m', 6)])
        node5 = self.create_raw_node(1, [])  # HIS
        node6 = self.create_raw_node(1, [])  # HIM
        node7 = self.create_raw_node(0, [('t', 8)])
        node8 = self.create_raw_node(1, [])  # IT

        self.count = 9
        self.raw = node0 + node1 + node2 + node3 + node4 + node5 + node6 + node7 + node8
        self.kind = ahocorasick.TRIE
        self.values = values
        self.word_count = 5

        A = self.create_automaton()
        self.assertEqual(len(A), 5)
        self.assertEqual(A.get("he"), "HE")
        self.assertEqual(A.get("her"), "HER")
        self.assertEqual(A.get("him"), "HIM")
        self.assertEqual(A.get("his"), "HIS")
        self.assertEqual(A.get("it"), "IT")

    @skipIf(not ahocorasick.unicode, "Run only with unicode build")
    def test__construct_simple_trie__split_across_a_few_chunks_unicode(self):

        r"""
        trie for set {he, her, his, him, it}

        #0 -> [h #1 ] -> [e #2*] -> [r #3*]
         |           \-> [i #4 ] -> [s #5*]
         |                      \-> [m #6*]
         |
         +--> [i #7 ] -> [t #8 ]
        """
        values = ["HE", "HER", "HIS", "HIM", "IT"]

        node0 = self.create_raw_node(0, [('h', 1), ('i', 7)])
        node1 = self.create_raw_node(0, [('e', 2), ('i', 4)])
        node2 = self.create_raw_node(1, [('r', 3)])  # HE
        node3 = self.create_raw_node(1, [])  # HER
        node4 = self.create_raw_node(0, [('s', 5), ('m', 6)])
        node5 = self.create_raw_node(1, [])  # HIS
        node6 = self.create_raw_node(1, [])  # HIM
        node7 = self.create_raw_node(0, [('t', 8)])
        node8 = self.create_raw_node(1, [])  # IT

        self.count = 9
        self.raw = [
            self.create_raw_count(2) + node0 + node1,
            self.create_raw_count(3) + node2 + node3 + node4,
            self.create_raw_count(1) + node5,
            self.create_raw_count(3) + node6 + node7 + node8
        ]
        self.kind = ahocorasick.TRIE
        self.values = values
        self.word_count = 5

        A = self.create_automaton(USE_EXACT_RAW)
        self.assertEqual(len(A), 5)
        self.assertEqual(A.get("he"), "HE")
        self.assertEqual(A.get("her"), "HER")
        self.assertEqual(A.get("him"), "HIM")
        self.assertEqual(A.get("his"), "HIS")
        self.assertEqual(A.get("it"), "IT")

    def test__construct_wrong_kind(self):

        self.kind = 10000

        with self.assertRaisesRegex(ValueError, "kind value.*"):
            self.create_automaton()

    def test__construct_wrong_store(self):

        self.store = 10000

        with self.assertRaisesRegex(ValueError, "store value.*"):
            self.create_automaton()

    def test__construct_wrong_key_type(self):

        self.key_type = 10000

        with self.assertRaisesRegex(ValueError, "key_type must have.*"):
            self.create_automaton()

    @skipIf(not ahocorasick.unicode, "Run only with unicode build")
    def test__construct_simple_trie__wrong_index_unicode(self):
        """
        trie for set {he}

        #0 -> [h #1*] -> [e #2*]
        """

        node0 = self.create_raw_node(0, [('h', 1)])
        node1 = self.create_raw_node(1, [('e', 2)])  # expect python value
        node2 = self.create_raw_node(1, [])  # also python value

        self.count = 3
        self.raw = node0 + node1 + node2
        self.kind = ahocorasick.TRIE
        self.values = ["HE"]  # but we provide a too short collection
        self.word_count = 2

        with self.assertRaises(IndexError):
            self.create_automaton()

    @skipIf(ahocorasick.unicode, "Run only with bytes build")
    def test__construct_simple_trie__wrong_index_bytes(self):
        """
        trie for set {he}

        #0 -> [h #1*] -> [e #2*]
        """

        node0 = self.create_raw_node(0, [(b'h', 1)])
        node1 = self.create_raw_node(1, [(b'e', 2)])  # expect python value
        node2 = self.create_raw_node(1, [])  # also python value

        self.count = 3
        self.raw = node0 + node1 + node2
        self.kind = ahocorasick.TRIE
        self.values = [b"HE"]  # but we provide a too short collection
        self.word_count = 2

        with self.assertRaises(ValueError):
            self.create_automaton()

    def test__truncated_raw__case_1(self):

        self.count = 1  # we're saying this is a non-empty trie, but given empty data
        self.raw = b''
        self.kind = ahocorasick.TRIE

        with self.assertRaisesRegex(ValueError, r"Data truncated \[parsing header of node #0\].*"):
            self.create_automaton()

    def test__truncated_raw__case_2(self):
        """
        trie for set {he}

        #0 -> [h #1 ] -> [e #2*]
        """

        node0 = self.create_raw_node(0, [('h', 1)])
        node1 = self.create_raw_node(0, [('e', 2)])
        node2 = self.create_raw_node(1, [])
        raw = node0 + node1 + node2

        self.count = 3
        self.kind = ahocorasick.TRIE

        for length in range(len(raw)):
            self.raw = raw[:length]  # truncate data and expect fail
            with self.assertRaisesRegex(ValueError, "Data truncated.*"):
                self.create_automaton()

    @skipIf(not ahocorasick.unicode, "Run only with unicode build")
    def test__malicious_next_pointer_unicode(self):
        """
        #0 -> [? #1 ]
        """

        node0 = self.create_raw_node(0, [('?', 1)])
        node1 = self.create_raw_node(0, [('x', 16)])  # the second node point to non-existent node

        self.count = 2
        self.raw = node0 + node1
        self.kind = ahocorasick.TRIE

        with self.assertRaisesRegex(ValueError, "Node #1 malformed: next link #0 points to.*"):
            self.create_automaton()

    @skipIf(ahocorasick.unicode, "Run only with bytes build")
    def test__malicious_next_pointer_bytes(self):
        """
        #0 -> [? #1 ]
        """

        node0 = self.create_raw_node(0, [('?', 1)])
        node1 = self.create_raw_node(0, [('x', 16)])  # the second node point to non-existent node

        self.count = 2
        self.raw = node0 + node1
        self.kind = ahocorasick.TRIE

        with self.assertRaisesRegex(ValueError, "Data truncated \\[parsing children of node #1\\]: chunk #0 @ offset.*"):
            self.create_automaton()

    def test__malicious_fail_pointer(self):
        """
        trie with just one node
        """

        builder = self.create_node_builder(0, [])
        builder.fail = 42

        self.count = 1
        self.raw = builder.dump()
        self.kind = ahocorasick.TRIE

        with self.assertRaisesRegex(ValueError, "Node #0 malformed: the fail link points to.*"):
            self.create_automaton()

    def test__values_leaks(self):

        # create not connected nodes, but each hold a value
        good_nodes = 1000
        raw = b''
        values = []
        for i in range(good_nodes):
            raw += self.create_raw_node(1, [])
            values.append(tuple("node %d" % i))

        # create the last node that will cause error -- malformed next pointer
        raw += self.create_raw_node(1, [('_', 10000)])
        values.append(tuple("never reached"))

        self.count = good_nodes + 1
        self.raw = raw
        self.kind = ahocorasick.TRIE
        self.values = values

        with self.assertRaises(ValueError):
            self.create_automaton()

    def test__wrong_type_of_chunk_container(self):

        self.count = 9
        self.raw = ()  # this should be a list
        self.kind = ahocorasick.TRIE
        self.values = None
        self.word_count = 5

        with self.assertRaisesRegex(TypeError, "Expected list"):
            A = self.create_automaton(USE_EXACT_RAW)

    def test__wrong_type_of_chunk(self):

        self.count = 9
        self.raw = [42]  # list items must be strings/bytes
        self.kind = ahocorasick.TRIE
        self.values = None
        self.word_count = 5

        with self.assertRaisesRegex(ValueError, "Item #0 on the bytes list is not a bytes object"):
            A = self.create_automaton(USE_EXACT_RAW)

    def test__wrong_count_of_nodes_in_chunk__case1(self):

        self.count = 9
        self.raw = [
            self.create_raw_count(0)  # count must be greater than 0
        ]
        self.kind = ahocorasick.TRIE
        self.values = None
        self.word_count = 5

        with self.assertRaisesRegex(ValueError, r"Nodes count for item #0 on the bytes list is not positive \(0\)"):
            A = self.create_automaton(USE_EXACT_RAW)

    def test__wrong_count_of_nodes_in_chunk__case2(self):

        self.count = 9
        self.raw = [
            self.create_raw_count(-12 & 0xffffffffffffffff)  # count must be greater than 0
        ]
        self.kind = ahocorasick.TRIE
        self.values = None
        self.word_count = 5

        try:
            A = self.create_automaton(USE_EXACT_RAW)
        except ValueError as e:
            assert str(e).startswith("Nodes count for item #0 on the bytes list is not positive")


if __name__ == '__main__':
    print("WARNING: these tests deal with in-memory representation (see TreeNodeBuilder),")
    print("         they were meant to test low-level implementation of pickling.")
    print("         Might segfault on your machine which is not necessary a bug in pyahocorasick.")
    unittest.main()
