#!/usr/bin/env python3

#
# Copyright (C) 2014 Noralf Tronnes
#
# MIT License
#

import io
import os
import struct
import sys
import argparse
import urllib.request, urllib.parse, urllib.error
import urllib.request, urllib.error, urllib.parse
import urllib.parse
import gzip
import subprocess
import re
import json
import platform
import errno

# used by archive file and unpacked archive
DISK_USAGE_MB = 900

PROCESSOR_TYPES_NAMES = ('BCM2835', 'BCM2836', 'BCM2837', 'BCM2711', 'BCM2712')
PROCESSOR_TYPES = range(0, len(PROCESSOR_TYPES_NAMES))
ARCHITECTURE_TYPES_NAMES = ('32-bit', '64-bit')
ARCHITECTURE_TYPES = range(0, len(ARCHITECTURE_TYPES_NAMES))

help = "https://github.com/RPi-Distro/rpi-source/blob/master/README.md"
script_repo = "https://github.com/RPi-Distro/rpi-source"

update_tag_file = os.path.join(os.environ.get('HOME'), '.rpi-source')
argv = sys.argv[:]

parser = argparse.ArgumentParser(description='Raspberry Pi kernel source installer',
                                  epilog="For more help see: %s" % help)
parser.add_argument("-d", "--dest", help="Destination directory. Default is $HOME",
                    default=os.environ.get('HOME'))
parser.add_argument("--nomake", help="Don't run 'make modules_prepare'",
                    action="store_true")
if os.environ.get('REPO_URI'):
    repo_uri = os.environ.get('REPO_URI')
else:
    repo_uri = "https://github.com/Hexxeh/rpi-firmware"
parser.add_argument("--uri", help="Github repository to use. Default is '%s'" % repo_uri,
                    default=repo_uri)
parser.add_argument("--delete", help="Delete downloaded archive",
                    action="store_true")
parser.add_argument("-s", "--dry-run", help="No action; perform a simulation of events that would occur but do not actually change the system.",
                    action="store_true")
parser.add_argument("-v", "--verbose", help="Verbose",
                    action="store_true")
parser.add_argument("-q", "--quiet", help="Quiet",
                    action="store_true")
parser.add_argument("-g", "--default-config", help="Generate a default kernel configuration with 'make bcmrpi_defconfig'",
                    action="store_true")
parser.add_argument("--processor", type=int, choices=PROCESSOR_TYPES, help="Override Processor type")
parser.add_argument("--architecture", type=int, choices=ARCHITECTURE_TYPES, help="Override Architecture type")
parser.add_argument("--skip-gcc", help=argparse.SUPPRESS, action="store_true")  # Deprecated
parser.add_argument("--skip-space", help="Skip disk space check",
                    action="store_true")
parser.add_argument("--skip-update", help="Skip checking for update to this script",
                    action="store_true")
parser.add_argument("--tag-update", help="Tell the update mechanism that this is the latest version of the script",
                    action="store_true")
parser.add_argument("--download-only", help="Just download the kernel tarball, without unpacking source and installing it",
                    action="store_true",)
args = parser.parse_args()


class Kernel:
    pass

def debug(str):
    if args.verbose:
        print(str)

def info(str):
    if not args.quiet:
        print("\n *** %s" % str)

def warn(str):
    if not args.quiet:
        print("\n !!! %s" % str)

def fail(str):
    sys.stderr.write("ERROR:\n%s\n\nHelp: %s\n" % (str, help))
    exit(1)

def sh(cmd):
    debug("%s" % cmd)
    if not args.dry_run:
        subprocess.check_call(cmd, shell=True)

def sh_out(cmd):
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = process.communicate()
    errcode = process.returncode
    if errcode:
        return None
    return out.decode('utf-8')

def writef(f, str, mode='w'):
    debug("writef(%s)" % f)
    if not args.dry_run:
        with open(f, mode) as f:
            f.write(str)

def download(url):
    debug("download: %s" % url)
    try:
        res = urllib.request.urlopen(url).read().decode('utf-8')
    except urllib.error.HTTPError as e:
        fail(
            "Couldn't download %s, HTTPError: %s\n\n%s"
            % (url, e.code, json.dumps(json.load(e), indent=4))
        )
    except urllib.error.URLError as e:
        fail("Couldn't download %s, URLError: %s" % (url, e.args))
    return res

def download_to(url, file):
    debug("download_to: %s -> %s" % (url, file))
    if not args.dry_run:
        urllib.request.urlretrieve (url, file)

def update_get_head():
    if update_get_head.ref:
        return update_get_head.ref
    repo_short = urllib.parse.urlparse(script_repo).path
    repo_api = "https://api.github.com/repos%s/git/refs/heads/master" % repo_short
    res = download(repo_api)
    try:
        j = json.loads(res)
    except ValueError:
        j = {}
    if 'object' in j and 'sha' in j['object']:
        update_get_head.ref = j['object']['sha']
        return update_get_head.ref
    else:
        warn("Self update: Could not get ref of last commit")
        debug("Github returned:\n%s" % res)
        return None
update_get_head.ref = None

def is_update_needed():
    debug("Check for update to rpi-source")
    if not os.path.exists(update_tag_file):
        return True
    ref = update_get_head()
    if not ref:
        return False
    with open(update_tag_file) as f:
        tag_file_ref = f.read().strip()
    if ref != tag_file_ref:
        return True
    else:
        return False

def update_tag():
    ref = update_get_head()
    if not ref:
       exit(1)
    info("Set update tag: %s" % ref)
    writef(update_tag_file, ref)

def do_update():
    # MOD: replace current script; do not assume script is located at /usr/bin/rpi-source; also keep ownership and permissions
    script_name = argv[0]
    info("Updating rpi-source")
    sh("sudo wget %s https://raw.githubusercontent.com/RPi-Distro/rpi-source/master/rpi-source -O %s" % ("-q" if args.quiet else "", script_name))
    update_tag()
    info("Restarting rpi-source")
    argv.insert(0, sys.executable)
    os.execv(sys.executable, argv)

def check_diskspace(dir):
    df = sh_out("df %s" % dir)
    nums = re.findall(r'(?<=\s)\d+(?=\s+)', df)
    if not nums or len(nums) != 3:
        info("Warning: unable to check available diskspace")
    if (int(nums[2]) / 1024) < DISK_USAGE_MB:
        fail("Not enough diskspace (%dMB) on %s\nSkip this check with --skip-space" % (DISK_USAGE_MB, dir))

# see if gcc major.minor version matches the one used to build the running kernel
def check_gcc():
    cmd = 'gcc -dumpversion'
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = process.communicate()
    errcode = process.returncode
    if errcode:
        debug("gcc version check failed: '%s' returned %d" %(cmd, errcode))
    gcc_ver = out.strip().decode('utf-8')

    with open('/proc/version', 'r') as f:
        proc_version = f.read()
    gcc_ver_kernel = re.search(r'gcc version (\d\.\d\.\d)', proc_version)

    if not gcc_ver or not gcc_ver_kernel:
        debug("gcc version check failed: could not extract version numbers")
        return
    a = gcc_ver.split('.')
    b = gcc_ver_kernel.group(1).split('.')

    if a[:2] == b[:2]:
        debug("gcc version check: OK")
    else:
        debug("gcc version check: mismatch between installed gcc (%s) and /proc/version (%s)" % (gcc_ver[0], gcc_ver_kernel.group(1)))

def proc_config_gz():
    if not os.path.exists('/proc/config.gz'):
        sh("sudo modprobe configs 2> /dev/null")

    if not os.path.exists('/proc/config.gz'):
        return ''

    with gzip.open('/proc/config.gz', 'rb') as f:
        return f.read().decode('utf-8')

def kernel_suffix():
    if processor_type not in PROCESSOR_TYPES:
        fail("Unsupported processor_type %d" % processor_type)

    if architecture_type not in ARCHITECTURE_TYPES:
        fail("Unsupported architecture_type %d" % architecture_type)

    if processor_type == 0: # BCM2835 uses kernel.img
        if architecture_type == 0:
            return ''
    elif processor_type == 1: # BCM2836 uses kernel7.img
        if architecture_type == 0:
            return '7'
    elif processor_type == 2: # BCM2837 uses kernel7.img or kernel8.img
        if architecture_type == 0:
            return '7'
        elif architecture_type == 1:
            return '8'
    elif processor_type == 3: # BCM2711 uses kernel7l.img or kernel8.img
        if architecture_type == 0:
            return '7l'
        elif architecture_type == 1:
            return '8'
    elif processor_type == 4: # BCM2712 uses kernel8.img or kernel_2712.img
        if architecture_type == 1 and page_size == 4096:
            return '8'
        if architecture_type == 1 and page_size == 16384:
            return '_2712'

    fail("Unsupported combination of processor_type %d and architecture_type %d" % (processor_type, architecture_type))

def rpi_update_method(uri):
    kernel = Kernel()

    info("rpi-update: %s" % uri)

    with open(firmware_revision_path) as f:
        fw_rev = f.read().strip()
    info("Firmware revision: %s" % fw_rev)

    repo_short = urllib.parse.urlparse(uri).path

    repo_api = "https://api.github.com/repos%s" % repo_short
    repo_raw = "https://raw.githubusercontent.com%s" % repo_short

    kernel.git_hash = download("%s/%s/git_hash" % (repo_raw, fw_rev)).strip()
    kernel.symvers = "%s/%s/Module%s.symvers" % (repo_raw, fw_rev, kernel_suffix())

    if not args.default_config:
        kernel.config = proc_config_gz()

    return kernel

def debian_method(fn):
    kernel = Kernel()
    info("Using: %s" % fn)
    with gzip.open(fn, 'rb') as f:
        debian_changelog = f.read()

    # Find first firmware entry in log (latest entries are at the top)
    fw_rev = re.search(r'firmware as of ([0-9a-fA-F]+)', debian_changelog.decode('utf-8'))
    if not fw_rev:
        fail("Could not identify latest firmware revision")

    fw_rev = fw_rev.group(1)

    info("Latest firmware revision: %s" % fw_rev)

    repo_raw = "https://raw.githubusercontent.com/raspberrypi/firmware"

    kernel.git_hash = download("%s/%s/extra/git_hash" % (repo_raw, fw_rev)).strip()
    kernel.symvers = "%s/%s/extra/Module%s.symvers" % (repo_raw, fw_rev, kernel_suffix())

    if not args.default_config:
        kernel.config = proc_config_gz()

    return kernel


# Taken from: https://github.com/gpiozero/gpiozero/blob/master/gpiozero/pins/local.py
# Copyright (c) 2016-2019 Dave Jones <dave@waveform.org.uk>
# Copyright (c) 2018 Martchus <martchus@gmx.net>
#
def get_revision():
    revision = None
    try:
        with io.open('/proc/device-tree/system/linux,revision', 'rb') as f:
            revision = hex(struct.unpack('>L', f.read(4))[0])[2:]
    except IOError as e:
        if e.errno != errno.ENOENT:
            raise e
        with io.open('/proc/cpuinfo', 'r') as f:
            for line in f:
                if line.startswith('Revision'):
                    revision = line.split(':')[1].strip().lower()
    if revision is None:
        fail("Unable to find board revision (use --processor argument)")
    try:
        return int(revision, base=16)
    except:
        fail("Unrecognized revision %r (use --processor argument)" % revision)


def get_processor_type():
    if args.processor is not None:
        return args.processor

    # The old boards don't have the detailed revision, so to keep things simple:
    if platform.machine() == "armv6l":
        return 0

    revision = get_revision()
    if not (revision & 0x800000):
        fail("Unexpected revision 0x%x (use --processor argument)" % revision)

    processor = (revision & 0xf000) >> 12
    if processor not in PROCESSOR_TYPES:
        fail("Unexpected processor %d (use --processor argument)" % processor)

    return processor


def get_architecture_type():
    if args.architecture is not None:
        return args.architecture

    if platform.machine() == 'aarch64':
        return 1
    else:
        return 0


def get_page_size():
    page_size = int(sh_out("getconf PAGESIZE"))
    if (page_size % 4096) != 0:
        fail("Unexpected page size: %d, exiting" % (page_size))
    return page_size


def check_bc():
    if not os.path.exists('/usr/bin/bc'):
        fail("bc is NOT installed. Needed by 'make modules_prepare'. On Raspberry Pi OS,\nrun 'sudo apt install bc' to install it.")

##############################################################################

processor_type = get_processor_type()
info("SoC: %s" % PROCESSOR_TYPES_NAMES[processor_type])

architecture_type = get_architecture_type()
info("Arch: %s" % ARCHITECTURE_TYPES_NAMES[architecture_type])

page_size = get_page_size()
info("Page Size: %d" % page_size)

# FIX usage of -d DEST with relative pathnames
args.dest = os.path.abspath(args.dest)

if os.path.exists('/boot/firmware/.firmware_revision'):
    firmware_revision_path = '/boot/firmware/.firmware_revision'
else:
    firmware_revision_path = '/boot/.firmware_revision'


if args.tag_update:
    update_tag()
    exit(0)

if not args.skip_update and is_update_needed():
    do_update()

if not os.path.isdir(args.dest):
    fail("Destination directory missing: %s" % args.dest)

if not args.skip_gcc:
    check_gcc()

check_bc()

debianlog = "/usr/share/doc/raspberrypi-bootloader/changelog.Debian.gz"
if os.path.exists(firmware_revision_path):
    kernel = rpi_update_method(args.uri)
elif os.path.exists(debianlog):
    kernel = debian_method(debianlog)
else:
    fail("Can't find a source for this kernel")

info("Linux source commit: %s" % kernel.git_hash)

linux_dir = os.path.join(args.dest, "linux-%s" % kernel.git_hash)
if os.path.exists(linux_dir):
    info("Kernel source already installed: %s\n" % linux_dir)
    exit(1)

if not args.skip_space:
    check_diskspace(args.dest)

linux_tar = os.path.join(args.dest, "linux-%s.tar.gz" % kernel.git_hash)
if not os.path.exists(linux_tar):
    info("Download kernel source")
    sh("wget %s -O %s https://github.com/raspberrypi/linux/archive/%s.tar.gz" % (("-q" if args.quiet else ""), linux_tar, kernel.git_hash))
else:
    info("Download kernel source: Already downloaded %s" % linux_tar)

if args.download_only:
    info("Downloaded kernel source tarball: %s" % linux_tar)
    quit()

info("Unpack kernel source")
if args.quiet:
    sh("cd %s && tar -xzf %s" % (args.dest, linux_tar))
else:
    sh("cd %s && tar --checkpoint=100 --checkpoint-action=dot -xzf %s" % (args.dest, linux_tar))

# This happens automatically in a git repo, but we use a tarball
info("Add '+' to kernel release string")
if not args.dry_run:
    with open(os.path.join(linux_dir, '.scmversion'),"w") as f:
        f.write("+")

linux_symlink = os.path.join(args.dest, 'linux')
info("Create symlink: %s" % linux_symlink)
sh("rm -f %s" % linux_symlink)
sh("ln -s %s %s" % (linux_dir, linux_symlink))

info("Create /lib/modules/<ver>/{build,source} symlinks")
sh("sudo rm -rf /lib/modules/$(uname -r)/build /lib/modules/$(uname -r)/source")
sh("sudo ln -sf %s /lib/modules/$(uname -r)/build" % linux_symlink)
sh("sudo ln -sf %s /lib/modules/$(uname -r)/source" % linux_symlink)

if args.default_config or not kernel.config:
    info(".config (generating default)")
    defconfig = None
    if processor_type == 0:
        defconfig = "bcmrpi_defconfig"
    elif processor_type == 1:
        defconfig = "bcm2709_defconfig"
    elif processor_type == 2:
        if architecture_type == 0:
            defconfig = "bcm2709_defconfig"
        elif architecture_type == 1:
            defconfig = "bcm2711_defconfig"
    elif processor_type == 3:
        defconfig = "bcm2711_defconfig"
    elif processor_type == 4:
        if page_size == 4096:
            defconfig = "bcm2711_defconfig"
        elif page_size == 16384:
            defconfig = "bcm2712_defconfig"
    if defconfig is not None:
        sh("cd %s && make %s" % (linux_symlink, defconfig))
    else:
        fail("Unsupported processor_type: %d, architecture_type: %d" % (processor_type, architecture_type))
else:
    info(".config")
    writef(os.path.join(linux_dir, '.config'), kernel.config)

info("Module.symvers")
download_to(kernel.symvers, os.path.join(linux_dir, "Module.symvers"))
sh("cd %s && cp -a Module.symvers Module.symvers.backup" % linux_dir)

if not args.nomake:
    info("make modules_prepare")
    sh("cd %s && make modules_prepare %s" % (linux_symlink, (" > /dev/null" if args.quiet else "")))

if not os.path.exists('/usr/include/ncurses.h'):
    info("ncurses-devel is NOT installed. Needed by 'make menuconfig'. On Raspberry Pi OS,\nrun 'sudo apt install libncurses5-dev' to install it.")

if args.delete:
    info("Delete downloaded archive")
    sh("rm %s" % linux_tar)

info("Help: %s" % help)