def test_insert_branching():
"""Test if insertion works correctly when branching."""
test_tree = trie.TrieTree()
test_tree.insert('fir')
test_tree.insert('fae')
assert len(test_tree.root.next_list) == 1
assert len(test_tree.root.next_list[0].next_list) == 2
def test_contains_not_in_tree():
"""Test if contains work correctly for string not in tree."""
test_tree = trie.TrieTree()
test_tree.insert('fir')
test_tree.insert('fae')
test_tree.insert('fox')
assert test_tree.contains('forest') is False
def main(*args):
try:
checkid = args[0][1]
if_use_dao = args[0][2]
except:
print("参数错误， 请填写 $checkid 和 $使用数据库数据（1）或者自定义和数据（0）")
return -100, "参数错误"
filterInputRule, filterOutputRule, natPreRule = util.GetRuleData(
checkid, if_use_dao) # 由filter表所有链上的规则组合而成
interface_set = util.GetInterfaceData(checkid) # nat表的preRouting链上的规则,网口数据
# 以下是测试数据
# rule0 = "Accept"
# rawRule1 = "ACCEPT tcp -- 192.168.1.1 0.0.0.0/0 tcp dpts:20000:30000\n"
# rawRule2 = "ACCEPT tcp -- 192.168.1.1/24 10.2.2.2 tcp dpts:22222\n"
# rawRule3 = "ACCEPT tcp -- 192.168.1.1/23 10.2.2.2/13 \n"
# rawRule4 = "ACCEPT tcp -- 192.168.1.1/20 10.168.2.2/12 tcp dpts:22:80\n"
# rawRule5 = "ACCEPT tcp -- 192.168.1.1/18 10.2.2.2/11 tcp dpts:22:80\n"
# rule_set.append(rule().rawRuleFormat(1, rawRule1))
# rule_set.append(rule().rawRuleFormat(2, rawRule2))
# rule_set.append(rule().rawRuleFormat(3, rawRule3))
# rule_set.append(rule().rawRuleFormat(4, rawRule4))
# rule_set.append(rule().rawRuleFormat(5, rawRule5))
# rawNatRule1 = "DNAT tcp -- 0.0.0.0/0 127.0.0.1 tcp dpt:7410 to:127.0.0.1:9200"
# rawNatRule2 = "DNAT tcp -- 0.0.0.0/0 10.2.2.2 tcp dpt:7410 to:10.2.2.2:9200"
# natPreRuleSet.append(rule().rawRuleFormat(1, rawNatRule1))
# natPreRuleSet.append(rule().rawRuleFormat(2, rawNatRule2))
# info0 = Interface.InterfaceInfo("ens33", "10.2.2.2", "255.255.255.0")
# info1 = Interface.InterfaceInfo("lo", "127.0.0.1", "0.0")
# interface_set.append(info0)
# interface_set.append(info1)
# 固定写法，生成针对源ip和目的ip的两颗分类树
sttree = trie.TrieTree(0)
dttree = trie.TrieTree(1)
# 规则集的冲突分析
ConflictAnalyse(filterInputRule, sttree, dttree)
# 结果输出
util.updateConflictAnalyseData(filterInputRule, 1, checkid)
# 多网口漏洞分析
Interface.InterfaceAnalyse(interface_set, dttree, natPreRule)
# 结果输出
util.updateNICAnalyseData(interface_set, checkid)
def test_remove_in_tree():
"""Test if contains work correctly for removing string in tree."""
test_tree = trie.TrieTree()
test_tree.insert('fir')
test_tree.insert('fae')
test_tree.insert('fox')
assert test_tree.contains('fox') is True
test_tree.remove('fox')
assert test_tree.contains('fox') is False
def test_remove_not_in_tree():
"""Test remove works correctly when word is not in tree."""
test_tree = trie.TrieTree()
test_tree.insert('fir')
test_tree.insert('faerie')
test_tree.insert('fox')
with pytest.raises(IndexError):
test_tree.remove('fae')
with pytest.raises(IndexError):
test_tree.remove('forest')
def test_traversal_from_single_letter():
"""Test if traversal of a single letters works correctly."""
test_tree = trie.TrieTree()
test_list = ['fae', 'fir', 'faerie', 'fox', 'forest']
for i in test_list:
test_tree.insert(i)
traversal_list = [i for i in test_tree.traversal('f')]
for i in test_list:
assert i in traversal_list
assert len(test_list) == len(traversal_list)
def test_size():
"""Test if size works correctly."""
test_tree = trie.TrieTree()
assert test_tree.size() == 0
test_tree.insert('fir')
test_tree.insert('faerie')
test_tree.insert('fox')
assert test_tree.size() == 3
test_tree.remove('fox')
assert test_tree.size() == 2
def test_insert_single_word(data, results):
"""Test if insertion works correctly for a single word."""
test_tree = trie.TrieTree()
test_tree.insert(data)
assert test_tree.root.next_list[0].value == results[0]
assert test_tree.root.next_list[0].end is False
assert test_tree.root.next_list[0].next_list[0].value == results[1]
assert test_tree.root.next_list[0].next_list[0].end is False
assert test_tree.root.next_list[0].next_list[0].next_list[
0].value == results[2]
assert test_tree.root.next_list[0].next_list[0].next_list[0].end is True
def test_remove_only_removes_endpoint_when_letters_have_children():
"""Test remove only removes endpoint when letters have children."""
test_tree = trie.TrieTree()
test_tree.insert('fir')
test_tree.insert('fae')
test_tree.insert('faerie')
test_tree.insert('fox')
assert test_tree.contains('fae') is True
test_tree.remove('fae')
assert test_tree.contains('fae') is False
assert test_tree.contains('faerie')
def test_trie_with_huge_database():
"""Import a gigantic dictionary and asserts that it works properly in trie tree."""
test_tree = trie.TrieTree()
with open('/usr/share/dict/words') as dictionary:
data = dictionary.read()
data = data.split('\n')
if len(data) > 100000:
data = data[:100000]
for i in range(len(data)):
test_tree.insert(data[i])
assert test_tree.contains('dinosaur')
test_tree.remove('dinosaur')
assert test_tree.contains('dinosaurs')
assert test_tree.contains('dinosaur') is False
def test_insert_empty_string():
"""Test if insertion works correctly with empty string."""
test_tree = trie.TrieTree()
with pytest.raises(IndexError):
test_tree.remove('')
def test_insert_invalid():
"""Test if insertion works correctly when invalid."""
test_tree = trie.TrieTree()
with pytest.raises(TypeError):
test_tree.remove(2)
def test_insert_duplicate():
"""Test if insertion works correctly for duplicate words."""
test_tree = trie.TrieTree()
test_tree.insert('twin')
test_tree.insert('twin')
assert len(test_tree.root.next_list) == 1