#!/usr/bin/env python3
# Unit tests for genseccomp.py

import textwrap
import unittest

import genseccomp

class TestGenseccomp(unittest.TestCase):
  def test_convert_NRs_to_ranges(self):
    ranges = genseccomp.convert_NRs_to_ranges([("b", 2), ("a", 1)])
    self.assertEqual(len(ranges), 1)
    self.assertEqual(ranges[0].begin, 1)
    self.assertEqual(ranges[0].end, 3)
    self.assertEqual(set(ranges[0].names), {"a", "b"})

    ranges = genseccomp.convert_NRs_to_ranges([("b", 3), ("a", 1)])
    self.assertEqual(len(ranges), 2)
    self.assertEqual(ranges[0].begin, 1)
    self.assertEqual(ranges[0].end, 2)
    self.assertEqual(set(ranges[0].names), {"a"})
    self.assertEqual(ranges[1].begin, 3)
    self.assertEqual(ranges[1].end, 4)
    self.assertEqual(set(ranges[1].names), {"b"})

  def test_convert_to_intermediate_bpf(self):
    ranges = genseccomp.convert_NRs_to_ranges([("b", 2), ("a", 1)])
    bpf = genseccomp.convert_to_intermediate_bpf(ranges)
    self.assertEqual(bpf, ['BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 3, {fail}, {allow}), //a|b'])

    ranges = genseccomp.convert_NRs_to_ranges([("b", 3), ("a", 1)])
    bpf = genseccomp.convert_to_intermediate_bpf(ranges)
    self.assertEqual(bpf, ['BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 3, 1, 0),',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 2, {fail}, {allow}), //a',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 4, {fail}, {allow}), //b'])

  def test_convert_ranges_to_bpf(self):
    ranges = genseccomp.convert_NRs_to_ranges([("b", 2), ("a", 1)])
    bpf = genseccomp.convert_ranges_to_bpf(ranges, priority_syscalls=[])
    self.assertEqual(bpf, ['BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 1, 0, 2),',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 3, 1, 0), //a|b',
                            'BPF_STMT(BPF_RET|BPF_K, SECCOMP_RET_ALLOW),'])

    ranges = genseccomp.convert_NRs_to_ranges([("b", 3), ("a", 1)])
    bpf = genseccomp.convert_ranges_to_bpf(ranges, priority_syscalls=[])
    self.assertEqual(bpf, ['BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 1, 0, 4),',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 3, 1, 0),',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 2, 2, 1), //a',
                            'BPF_JUMP(BPF_JMP|BPF_JGE|BPF_K, 4, 1, 0), //b',
                            'BPF_STMT(BPF_RET|BPF_K, SECCOMP_RET_ALLOW),'])

  def test_convert_bpf_to_output(self):
    output = genseccomp.convert_bpf_to_output(["line1", "line2"],
                                              "arm",
                                              name_modifier="")
    expected_output = textwrap.dedent("""\
    // File autogenerated by genseccomp.py - edit at your peril!!

    #include <linux/filter.h>
    #include <errno.h>

    #include "seccomp/seccomp_bpfs.h"
    const sock_filter arm_filter[] = {
    line1
    line2
    };

    const size_t arm_filter_size = sizeof(arm_filter) / sizeof(struct sock_filter);
    """)
    self.assertEqual(output, expected_output)


if __name__ == '__main__':
  unittest.main()