#!/usr/bin/env python3
# Copyright (C) 2020 Denis 'GNUtoo' Carikli <GNUtoo@cyberdimension.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
import re
import sh
import sys

def usage(progname):
    print("Usage:")
    print("    {} dd [path/to/mmcblk0.img] <path/to/file.pit> "
          "# Prints dd commands to extract the partitions".format(progname))
    print("    {} redmine <path/to/file.pit> "
          "# Prints a redmine table describing the partitions".format(progname))
    sys.exit(1)

def get_human_size(size):
    size_names = {
        1024*1024 : 'GiB',
        1024 : 'KiB',
    }

    if size == 0:
        return '0'

    elif size == None:
        return None

    for k, v in size_names.items():
        if (size / k) < 1:
            continue
        elif size % k == 0:
            return "{}{}".format((size / k), v)

    return size

def generate_dd_commands(pit_data, input_file='mmcblk0'):
    results = []

    block_size = 512

    for k, v in pit_data.items():
        output_file = v.get('partition_name', None)
        partition_location = v.get('partition_block_offset', None)
        partition_size = v.get('partition_block_count', None)

        if not output_file or not partition_location or not partition_size:
            continue

        results.append("# dd bs={} if={} of={} skip={} count={}".format(
            block_size,
            input_file,
            output_file,
            int (partition_location / 512),
            int (partition_size / 512),
        ))

    return os.linesep.join(results)

# TODO: also generate tables for mediawiki
def generate_redmine_table(pit_data):
    results = ""

    # Top header
    results += "|_. PIT "
    results += "|_. Linux name "
    results += "|_. Block device "
    results += "|_. Mount point "
    results += "|_. Partition type "
    results += "|._ Description "
    results += "|_. size or size@location for 16G eMMC "
    results += "|" + os.linesep

    for k, v in pit_data.items():
        results += "| {} ".format(v.get('partition_name', ''))
        results += "| "
        results += "| "
        results += "| "
        results += "| "
        results += "| "

        partition_size = get_human_size(v.get('partition_block_count', None))
        partition_location = get_human_size(
            v.get('partition_block_offset', None))

        results += "| "
        if partition_size:
            results += "{}".format(partition_size)
        if partition_location:
            results += "@{}".format(partition_location)
        results += "|" + os.linesep

    return results

# TODO:
# - hook libpit instead
# - Handle and decode the following information:
#   - Binary Type: 0 (AP)
#   - Device Type: 2 (MMC)
#   - Attributes: 5 (Read/Write)
#   - Update Attributes: 5 (FOTA)
def parse_pit(pit_file_path):
    results = {}
    entry_nr = None

    output = sh.heimdall('print-pit', '--file', pit_file_path)
    for line in output.split(os.linesep):
        found = re.search('^Entry Count: (.*)$', line)
        if found:
            nr_entries = int(found.group(1))
            continue

        found = re.search('^--- Entry #(.*) ---$', line)
        if found:
            entry_nr = int(found.group(1))
            results[str(entry_nr)] = {}
            continue

        found = re.search('^Partition Block Size/Offset: (.*)$', line)
        if found:
            partition_block_offset = int(found.group(1)) * 512
            results[str(entry_nr)]['partition_block_offset'] = \
                partition_block_offset
            continue

        found = re.search('^Partition Block Count: (.*)$', line)
        if found:
            partition_block_count = int(found.group(1)) * 512
            results[str(entry_nr)]['partition_block_count'] = \
                partition_block_count
            continue

        found = re.search('^Partition Name: (.*)$', line)
        if found:
            partition_name = str(found.group(1))
            results[str(entry_nr)]['partition_name'] = partition_name
            continue

        found = re.search('^Partition Name: (.*)$', line)
        if found:
            identifier = int(found.group(1))
            results[str(entry_nr)]['identifier'] = identifier
            continue

        found = re.search('^Flash Filename: (.*)$', line)
        if found:
            flash_filename = str(found.group(1))
            results[str(entry_nr)]['flash_filename'] = flash_filename
            continue

        found = re.search('^FOTA Filename: (.*)$', line)
        if found:
            fota_filename = str(found.group(1))
            results[str(entry_nr)]['fota_filename'] = fota_filename
            continue

        found = re.search('^File Offset \(Obsolete\): (.*)$', line)
        if found:
            obsolete_file_offset = str(found.group(1))
            results[str(entry_nr)]['obsolete_file_offset'] = \
                obsolete_file_offset
            continue

        found = re.search('^File Size \(Obsolete\): (.*)$', line)
        if found:
            obsolete_file_size = str(found.group(1))
            results[str(entry_nr)]['obsolete_file_size'] = \
                obsolete_file_offset
            continue

    return results


if __name__ == "__main__":
    if len(sys.argv) == 3 and sys.argv[1] == 'redmine':
        pit_data = parse_pit(sys.argv[2])
        print(generate_redmine_table(pit_data))
    elif len(sys.argv) == 3 and sys.argv[1] == 'dd':
        pit_data = parse_pit(sys.argv[2])
        print(generate_dd_commands(pit_data))
    elif len(sys.argv) == 4 and sys.argv[1] == 'dd':
        pit_data = parse_pit(sys.argv[2])
        mmcblk0_path = sys.argv[3]
        print(generate_dd_commands(pit_data, mmcblk0_path))
    else:
        usage(sys.argv[0])