1#!/usr/bin/env python3
2
3from __future__ import print_function
4
5import os
6import subprocess
7import unittest
8import zipfile
9
10from vndk_definition_tool import DexFileReader, UnicodeSurrogateDecodeError
11
12from .compat import TemporaryDirectory
13
14SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
15INPUT_DIR = os.path.join(SCRIPT_DIR, 'testdata', 'test_dex_file')
16
17
18class ModifiedUTF8Test(unittest.TestCase):
19    def test_encode(self):
20        self.assertEqual(b'\xc0\x80', u'\u0000'.encode('mutf-8'))
21        self.assertEqual(b'\x09', u'\u0009'.encode('mutf-8'))
22        self.assertEqual(b'\x7f', u'\u007f'.encode('mutf-8'))
23        self.assertEqual(b'\xc2\x80', u'\u0080'.encode('mutf-8'))
24        self.assertEqual(b'\xdf\xbf', u'\u07ff'.encode('mutf-8'))
25        self.assertEqual(b'\xe0\xa0\x80', u'\u0800'.encode('mutf-8'))
26        self.assertEqual(b'\xe7\xbf\xbf', u'\u7fff'.encode('mutf-8'))
27        self.assertEqual(b'\xed\xa0\x81\xed\xb0\x80',
28                         u'\U00010400'.encode('mutf-8'))
29
30
31    def test_decode(self):
32        self.assertEqual(u'\u0000', b'\xc0\x80'.decode('mutf-8'))
33        self.assertEqual(u'\u0009', b'\x09'.decode('mutf-8'))
34        self.assertEqual(u'\u007f', b'\x7f'.decode('mutf-8'))
35        self.assertEqual(u'\u0080', b'\xc2\x80'.decode('mutf-8'))
36        self.assertEqual(u'\u07ff', b'\xdf\xbf'.decode('mutf-8'))
37        self.assertEqual(u'\u0800', b'\xe0\xa0\x80'.decode('mutf-8'))
38        self.assertEqual(u'\u7fff', b'\xe7\xbf\xbf'.decode('mutf-8'))
39        self.assertEqual(u'\U00010400',
40                         b'\xed\xa0\x81\xed\xb0\x80'.decode('mutf-8'))
41
42
43    def test_decode_exception(self):
44        # Low surrogate does not come after high surrogate
45        with self.assertRaises(UnicodeSurrogateDecodeError):
46            b'\xed\xa0\x81\x40'.decode('mutf-8')
47
48        # Low surrogate without prior high surrogate
49        with self.assertRaises(UnicodeSurrogateDecodeError):
50            b'\xed\xb0\x80\x40'.decode('mutf-8')
51
52        # Unexpected end after high surrogate
53        with self.assertRaises(UnicodeSurrogateDecodeError):
54            b'\xed\xa0\x81'.decode('mutf-8')
55
56        # Unexpected end after low surrogate
57        with self.assertRaises(UnicodeSurrogateDecodeError):
58            b'\xed\xb0\x80'.decode('mutf-8')
59
60        # Out-of-order surrogate
61        with self.assertRaises(UnicodeSurrogateDecodeError):
62            b'\xed\xb0\x80\xed\xa0\x81'.decode('mutf-8')
63
64
65class DexFileTest(unittest.TestCase):
66    def _assemble_smali(self, dest, source):
67        """Assemble a smali source file.  Skip the test if the smali command is
68        unavailable."""
69        try:
70            subprocess.check_call(['smali', 'a', source, '-o', dest])
71        except EnvironmentError:
72            self.skipTest('smali not available')
73
74
75    def _create_zip_file(self, dest, paths):
76        """Create a zip file from several input files."""
77        with zipfile.ZipFile(dest, 'w') as zip_file:
78            for path in paths:
79                zip_file.write(path, os.path.basename(path))
80
81
82    def test_generate_classes_dex_names(self):
83        seq = DexFileReader.generate_classes_dex_names()
84        self.assertEqual('classes.dex', next(seq))
85        self.assertEqual('classes2.dex', next(seq))
86        self.assertEqual('classes3.dex', next(seq))
87
88
89    def test_enumerate_dex_strings_buf(self):
90        with TemporaryDirectory() as tmp_dir:
91            smali_file = os.path.join(INPUT_DIR, 'Hello.smali')
92            classes_dex = os.path.join(tmp_dir, 'classes.dex')
93            self._assemble_smali(classes_dex, smali_file)
94
95            with open(classes_dex, 'rb') as classes_dex:
96                buf = classes_dex.read()
97
98            strs = set(DexFileReader.enumerate_dex_strings_buf(buf))
99
100            self.assertIn(b'hello', strs)
101            self.assertIn(b'world', strs)
102
103
104    def test_enumerate_dex_strings_apk(self):
105        with TemporaryDirectory() as tmp_dir:
106            smali_file = os.path.join(INPUT_DIR, 'Hello.smali')
107            classes_dex = os.path.join(tmp_dir, 'classes.dex')
108            self._assemble_smali(classes_dex, smali_file)
109
110            smali_file = os.path.join(INPUT_DIR, 'Example.smali')
111            classes2_dex = os.path.join(tmp_dir, 'classes2.dex')
112            self._assemble_smali(classes2_dex, smali_file)
113
114            zip_file = os.path.join(tmp_dir, 'example.apk')
115            self._create_zip_file(zip_file, [classes_dex, classes2_dex])
116
117            strs = set(DexFileReader.enumerate_dex_strings_apk(zip_file))
118
119            self.assertIn(b'hello', strs)
120            self.assertIn(b'world', strs)
121            self.assertIn(b'foo', strs)
122            self.assertIn(b'bar', strs)
123