#!/usr/bin/env python
#
# Copyright 2017 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import subprocess
import tempfile
import unittest
import zipfile

import emmc_image


class GPTPartitionTableTest(unittest.TestCase):

  def testGPTPartitionTable(self):
    gpt_header = ('EFI PART'
                  '\x00\x00\x01\x00'
                  '\\\x00\x00\x00'
                  '\x88R\xe9\x92'
                  '\x00\x00\x00\x00'
                  '\x01\x00\x00\x00\x00\x00\x00\x00'
                  '\xff\xff\x7f\x00\x00\x00\x00\x00'
                  '"\x00\x00\x00\x00\x00\x00\x00'
                  '\xde\xff\x7f\x00\x00\x00\x00\x00'
                  '\xdc\xdd0\x7f-WVF\xa3\x81\x96W\xfd\xe5\\L'
                  '\x02\x00\x00\x00\x00\x00\x00\x00'
                  '\x80\x00\x00\x00'
                  '\x80\x00\x00\x00'
                  '\x93J\xf4]')
    gpt = emmc_image.GPTPartitionTable(gpt_header)
    self.assertEqual(gpt.signature, 'EFI PART')
    self.assertEqual(gpt.revision, 65536)
    self.assertEqual(gpt.header_size, len(gpt_header))
    self.assertEqual(gpt.crc32, 0x92e95288)
    self.assertEqual(gpt.current_lba, 1)
    self.assertEqual(gpt.backup_lba, 0x7fffff)
    self.assertEqual(gpt.first_usable_lba, 34)
    self.assertEqual(gpt.last_usable_lba, 0x7fffde)
    self.assertEqual(gpt.disk_guid, '\xdc\xdd0\x7f-WVF\xa3\x81\x96W\xfd\xe5\\L')
    self.assertEqual(gpt.part_entry_start_lba, 2)
    self.assertEqual(gpt.num_part_entries, 128)
    self.assertEqual(gpt.part_entry_size, 128)
    self.assertEqual(gpt.crc32_part_array, 0x5df44a93)


class GenerateEMMCTest(unittest.TestCase):

  def setUp(self):
    bpt = '''
    {
        "settings": {
            "disk_size": "1 MiB"
        },
        "partitions": [
            {
                "ab": true,
                "label": "boot",
                "size": "10 KiB"
            },
            {
                "ab": true,
                "label": "system",
                "size": "20 KiB"
            },
            {
                "label": "misc",
                "size": "1 KiB"
            }
        ]
    }'''

    bpt_file = tempfile.NamedTemporaryFile()
    bpt_file.write(bpt)
    bpt_file.flush()
    self.gpt_file = tempfile.NamedTemporaryFile()
    subprocess.check_call([
        'bpttool', 'make_table', '--input', bpt_file.name, '--output_gpt',
        self.gpt_file.name
    ])

    boot_file = tempfile.NamedTemporaryFile()
    self.boot = os.urandom(8 * 1024)
    boot_file.write(self.boot)
    boot_file.flush()

    zip_file = tempfile.NamedTemporaryFile()
    self.misc = os.urandom(300)
    with zipfile.ZipFile(zip_file.name, 'w') as factory_files_zip:
      factory_files_zip.writestr('misc.img', self.misc)
    self.out_file = tempfile.NamedTemporaryFile()

    class FakeArgs():

      def __init__(self):
        self.partition_table = ''
        self.partition = [['boot', boot_file.name]]
        self.factory_files = zip_file.name
        self.tfzips = []
        self.out = ''
        self.all_slots = False

    self.args = FakeArgs()
    self.args.partition_table = self.gpt_file.name
    self.args.out = self.out_file.name


  def testGenerateEMMC(self):
    self.assertEqual(emmc_image.GenerateEMMC(self.args), 0)

    out_zip = zipfile.ZipFile(self.out_file.name, 'r')
    emmc_file = out_zip.open('emmc.img', 'r')
    emmc = emmc_file.read()
    self.assertEqual(len(emmc), 1024 * 1024)

    self.assertEqual(emmc[:17 * 1024], self.gpt_file.read(17 * 1024))
    boot_offset = 20 * 1024
    boot_part_size = 10 * 1024
    self.assertEqual(emmc[boot_offset:boot_offset + len(self.boot)], self.boot)
    # rest of boot_a should be 0
    self.assertEqual(
        emmc[boot_offset + len(self.boot):boot_offset + boot_part_size],
        '\0' * (boot_part_size - len(self.boot)))
    # padding between boot_a and boot_b should be 0
    self.assertEqual(emmc[30 * 1024:32 * 1024], '\0' * 2 * 1024)
    # unwritten boot_b should be 0
    self.assertEqual(emmc[32 * 1024:42 * 1024], '\0' * 10 * 1024)
    # padding between boot_b and system_a should be 0
    self.assertEqual(emmc[30 * 1024:32 * 1024], '\0' * 2 * 1024)
    # unwritten system_a should be 0
    self.assertEqual(emmc[44 * 1024:64 * 1024], '\0' * 20 * 1024)
    # unwritten system_b should be 0
    self.assertEqual(emmc[64 * 1024:84 * 1024], '\0' * 20 * 1024)
    misc_offset = 84 * 1024
    misc_part_size = 1024
    self.assertEqual(emmc[misc_offset:misc_offset + len(self.misc)], self.misc)
    # rest of misc should be 0
    self.assertEqual(
        emmc[misc_offset + len(self.misc):misc_offset + misc_part_size],
        '\0' * (misc_part_size - len(self.misc)))
    # secondary gpt
    self.assertEqual(emmc[-33 * 512:], self.gpt_file.read(33 * 512))

  def testParitionOverride(self):
    boot_b_file = tempfile.NamedTemporaryFile()
    boot_b = os.urandom(9 * 1024)
    boot_b_file.write(boot_b)
    boot_b_file.flush()

    self.args.partition += [['boot_b', boot_b_file.name]]
    self.args.all_slots = True
    self.assertEqual(emmc_image.GenerateEMMC(self.args), 0)

    out_zip = zipfile.ZipFile(self.out_file.name, 'r')
    emmc_file = out_zip.open('emmc.img', 'r')
    emmc = emmc_file.read()
    self.assertEqual(len(emmc), 1024 * 1024)

    self.assertEqual(emmc[:17 * 1024], self.gpt_file.read(17 * 1024))
    boot_offset = 20 * 1024
    boot_part_size = 10 * 1024
    self.assertEqual(emmc[boot_offset:boot_offset + len(self.boot)], self.boot)
    # rest of boot_a should be 0
    self.assertEqual(
        emmc[boot_offset + len(self.boot):boot_offset + boot_part_size],
        '\0' * (boot_part_size - len(self.boot)))
    # padding between boot_a and boot_b should be 0
    self.assertEqual(emmc[30 * 1024:32 * 1024], '\0' * 2 * 1024)
    boot_b_offset = 32 * 1024
    self.assertEqual(emmc[boot_b_offset:boot_b_offset + len(boot_b)], boot_b)
    # rest of boot_b should be 0
    self.assertEqual(
        emmc[boot_b_offset + len(boot_b):boot_b_offset + boot_part_size],
        '\0' * (boot_part_size - len(boot_b)))
    # padding between boot_b and system_a should be 0
    self.assertEqual(emmc[30 * 1024:32 * 1024], '\0' * 2 * 1024)
    # unwritten system_a should be 0
    self.assertEqual(emmc[44 * 1024:64 * 1024], '\0' * 20 * 1024)
    # unwritten system_b should be 0
    self.assertEqual(emmc[64 * 1024:84 * 1024], '\0' * 20 * 1024)
    misc_offset = 84 * 1024
    misc_part_size = 1024
    self.assertEqual(emmc[misc_offset:misc_offset + len(self.misc)], self.misc)
    # rest of misc should be 0
    self.assertEqual(
        emmc[misc_offset + len(self.misc):misc_offset + misc_part_size],
        '\0' * (misc_part_size - len(self.misc)))
    # secondary gpt
    self.assertEqual(emmc[-33 * 512:], self.gpt_file.read(33 * 512))


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