flash_card.py

#!/bin/python3
import glob
import os
import random
import string
import shutil
import subprocess
import tarfile
import argparse
import io
import sys
from tempfile import TemporaryDirectory

__author__ = "Gareth Dunstone"
__copyright__ = "Copyright 2016, Borevitz Lab"
__credits__ = ["Gareth Dunstone", "Tim Brown", "Justin Borevitz", "Kevin Murray", "Jack Adamson"]
__license__ = "GPL"
__version__ = "3.0.2"
__maintainer__ = "Gareth Dunstone"
__email__ = "gareth.dunstone@anu.edu.au"
__status__ = "Testing"

parser = argparse.ArgumentParser(description="Flash spc-eyepi sd card")
parser.add_argument("blockdevice", metavar='d', type=str, nargs=1,
                    help="Block device/devices to copy data to")
parser.add_argument("--tarfile", metavar="t", type=argparse.FileType('rb'),
                    help="tar/tar.gz file to flash to the card")
parser.add_argument("--to-tarfile", metavar="w", type=str,
                    help="Copy from card")
parser.add_argument("--update-boot", metavar="w", type=str,
                    help="update raspberry pi boot partition from a directory")

parser.add_argument("--api-token", metavar='k', type=str,
                    help="traitcapture api token for automated addition to database")

parser.add_argument("--update", default=False, action='store_true',
                    help="dont flash new data to the card, update the software and set the name if required.")
parser.add_argument("--remove-keys", default=False, action='store_true',
                    help="clear ssh keys.")
parser.add_argument("--remove-tor", default=False, action='store_true',
                    help="clear tor private keys.")
parser.add_argument("--remove-configs", default=False, action='store_true',
                    help="clear previous configs.")
parser.add_argument("--reset-machine-id", default=False, action='store_true',
                    help="clear tor private keys.")
parser.add_argument("--backup", metavar='b',
                    help="backup the tor encryption keys, ssh encryption keys, and camera config files to a directory.")
parser.add_argument("--restore", metavar='r',
                    help="restore tor encryption keys, ssh encryption keys, and camera config files from a directory.")
parser.add_argument("--name", metavar='n', type=str,
                    help="name of the raspberry pi (will be autogenerated if not provided)")

command_line_args = parser.parse_args()


class BColors:
    """
    Colours for the terminal
    """
    header = '\033[95m'
    blue = '\033[94m'
    green = '\033[92m'
    warn = '\033[93m'
    fail = '\033[91m'
    endc = '\033[0m'
    bold = '\033[1m'
    under = '\033[4m'


def printc(text, *args):
    """
    prints colors
    :param text: text to print
    :param args: colours from BColors
    :return:
    """
    print("".join(args) + text + BColors.endc)


def printr(*text):
    """
    prints to the terminal with carriage return.
    :param text: strings to print
    :return:
    """
    # print("\x1b[1G\x1b[K", text, sep="", end="")
    print(*text, sep="", end="\r")


RAND = "".join((random.choice(string.ascii_letters) for _ in range(6)))
if not command_line_args.name:
    printc("No name provided. autogenerated: {} for a name".format("Picam-" + RAND), BColors.warn)
gname = "Picam-" + RAND if not command_line_args.name else command_line_args.name

global_char_pos = 0


def progressbar(current, total):
    """
    displays a progressbar
    :param current: current amount
    :param total: total amount
    :return:
    """
    global global_char_pos
    term_width = shutil.get_terminal_size((80, 20)).columns - 6
    progress_char_pos = int((current / total) * term_width)
    if progress_char_pos != global_char_pos:
        col = BColors.green
        if progress_char_pos <= term_width / 3:
            col = BColors.warn
        elif term_width / 3 <= progress_char_pos <= term_width * (2 / 3):
            col = BColors.header
        s = "".join("X" if x < progress_char_pos else " " for x in range(term_width))
        printr(col + "-=[{}]=-".format(s) + BColors.endc)
        global_char_pos = progress_char_pos


def write_key_token(tmpdir):
    """
    writes a key token from traitcapture.org
    also refreshes the token.
    :param tmpdir: temporary directory where sd card is mounted.
    :return:
    """
    import requests
    api_token = command_line_args.api_token
    if os.path.exists(command_line_args.api_token):
        with open(command_line_args.api_token, 'r') as f:
            api_token = f.read().strip()

    resp = requests.get("https://traitcapture.org/api/code/new/14.jsonp?token=" + api_token)
    if resp.status_code == 200:
        try:
            # js = json.loads(resp.text)
            js = resp.json()
            from pprint import pformat
            printc(pformat(js), BColors.blue)
            ssh_dir = os.path.join(tmpdir, "root", "home", ".ssh")
            os.makedirs(ssh_dir, exist_ok=True)
            with open(os.path.join(ssh_dir, "key_token"), 'w') as f:
                f.write(js['code'])

            if os.path.exists(command_line_args.api_token):
                with open(command_line_args.api_token, 'w') as f:
                    printc("Refreshing api token", BColors.blue)
                    response = requests.get("https://traitcapture.org/api/token.json?token=" + api_token)
                    if response.status_code == 200:
                        js = response.json()
                        f.write(js["token"])
                        printc("Token refreshed successfully", BColors.blue)
                    else:
                        printc("refreshing api token failed, might need to get a new one", BColors.fail)
        except KeyError:
            printc("Server didnt respond with a valid code", BColors.fail)

        except Exception as e:
            printc("Couldnt get token using key: {}".format(str(e)), BColors.fail)

    else:
        printc("invalid response from server {}".format(str(resp)), BColors.fail)


def backup(tmpdir):
    """
    backs up a cards important data, like ssh keys, tor keys, and configuration files.
    :param tmpdir:
    :return:
    """
    printc("copying old files over...", BColors.blue)
    with open(os.path.join(tmpdir, "root","etc", "hostname"), 'r') as f:
        hostname = f.read().strip()
    bakdir = "{}.bak".format(hostname)
    try:
        os.makedirs(bakdir)
    except:
        pass
    shutil.copy(os.path.join(tmpdir, "root", "etc", "hostname"), os.path.join(bakdir, "hostname"))
    shutil.copytree(os.path.join(tmpdir, "root", "home", "spc-eyepi", "configs_byserial"), os.path.join(bakdir, "configs"))
    shutil.copytree(os.path.join(tmpdir, "root", "home", "tor_private"), os.path.join(bakdir, "tor_private"))
    shutil.copytree(os.path.join(tmpdir, "root", "home", ".ssh"), os.path.join(bakdir, "ssh"))


def restore(tmpdir, bakdir=None):
    """
    the opposite of backup.
    restores critical information to the raspberry pi.
    :param tmpdir:
    :param bakdir: path to the backup to restore
    :return:
    """
    if not bakdir:
        with open(os.path.join(tmpdir, "root", "hostname"), 'r') as f:
            hostname = f.read().strip()
        bakdir = "{}.bak".format(hostname)

    if os.path.exists(bakdir):
        shutil.copy(os.path.join(bakdir, "hostname"), os.path.join(tmpdir, "root", "etc", "hostname"))
        config_files = glob.glob(os.path.join(bakdir, "configs", "*"))
        for f in config_files:
            shutil.copy(f, os.path.join(tmpdir, "root", "home", "spc-eyepi", "configs_byserial"))
        tor_files = glob.glob(os.path.join(bakdir, "tor_private", "*"))
        for f in tor_files:
            shutil.copy(f, os.path.join(tmpdir, "root", "home", "tor_private"))
        if not os.path.exists(os.path.join(tmpdir, "root", "home", ".ssh")):
            os.makedirs(os.path.join(tmpdir, "root", "home", ".ssh"))
        ssh_files = glob.glob(os.path.join(bakdir, "ssh","*"))
        for f in ssh_files:
            shutil.copy(f, os.path.join(tmpdir, "root", "home", "ssh"))

        if os.path.isfile(os.path.join(bakdir, "hostname")):
            with open(os.path.join(bakdir, "hostname"), 'r') as hostname_file:
                with open(os.path.join(tmpdir, "root", "etc", "hosts"), 'w') as hosts_file:
                    h_tmpl = "127.0.0.1\tlocalhost.localdomain localhost {hostname}\n"
                    h_tmpl += "::1\tlocalhost.localdomain localhost {hostname}\n"
                    hosts_file.write(h_tmpl.format(hostname=hostname_file.read().strip()))
        printc("Completed restore", BColors.blue)
    else:
        printc("Error: backup dir doesnt exist...", BColors.fail)


def update_via_github(tmpdir):
    """
    updates spc-eyepi on the card from github
    :param tmpdir:
    :return:
    """
    eyepi_dir = os.path.join(tmpdir, "root", "home", "spc-eyepi")
    git_dir = os.path.join(eyepi_dir, ".git")
    try:
        x = subprocess.run(["git --git-dir={} --work-tree={} fetch --all".format(git_dir, eyepi_dir)],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                           shell=True, universal_newlines=True)
        x.check_returncode()
        x = subprocess.run(['git --git-dir={} --work-tree={} reset --hard origin/master'.format(git_dir, eyepi_dir)],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                           shell=True, universal_newlines=True)
        x.check_returncode()

        v = subprocess.check_output(["git --git-dir={} --work-tree={} describe".format(git_dir, eyepi_dir)],
                                    shell=True, universal_newlines=True)
        q = subprocess.check_output(["git --git-dir={} --work-tree={} log -1 --pretty=%B".format(git_dir, eyepi_dir)],
                                    shell=True, universal_newlines=True)
        z = subprocess.check_output(
            ["git --git-dir={} --work-tree={} show -s --format=%cd --date=local".format(git_dir, eyepi_dir)],
            shell=True, universal_newlines=True)

        printc("Now at:\n\n{}{}{}".format(v, z, q), BColors.blue)
        printc("Completed update", BColors.blue)
    except subprocess.CalledProcessError as e:
        printc("Couldnt call git: {}".format(str(e)), BColors.fail)
    except Exception as e:
        printc("Exception while updating from github: {}".format(str(e)), BColors.fail)


def partition_and_format():
    """
    partitions and formats an sd card in preparation to have an os extracted to it.
    :return:
    """
    printc("Creating new partition table with fdisk...", BColors.warn)
    with subprocess.Popen(["/usr/bin/fdisk", "{}".format(command_line_args.blockdevice[0])],
                          stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                          universal_newlines=True) as proc:
        # create boot partition of size 100M
        proc.stdin.write("o\nn\np\n\n\n+100M\n")
        proc.stdin.flush()
        # make w95 fat for rpi :(
        proc.stdin.write("t\nc\n")
        proc.stdin.flush()
        # create linux partition to fill the rest of the disk.
        proc.stdin.write("n\n\n\n\n\n")
        proc.stdin.flush()

        # write to disk
        proc.stdin.write("w\n")
        proc.stdin.flush()
        # get rid of fdisk.
        proc.communicate()

    printc("Creating fat32 boot filesystem (for raspberry pi firmware and bootloader).", BColors.warn)
    with subprocess.Popen(["/usr/bin/mkfs.vfat", "{}".format(command_line_args.blockdevice[0] + "1")],
                          stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                          universal_newlines=True) as proc:
        proc.communicate()

    printc("Creating ext4  root filesystem (for the os).", BColors.warn)
    with subprocess.Popen(["/usr/bin/mkfs.ext4", "{}".format(command_line_args.blockdevice[0] + "2")],
                          stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                          universal_newlines=True) as proc:
        proc.communicate()


class ProgressFileObject(io.FileIO):
    """
    overriden file object to give tarfile the ability to progressbar.
    """

    def __init__(self, path, *args, **kwargs):
        self._total_size = max(os.path.getsize(path), 2666188800)
        super().__init__(path, *args, **kwargs)
        # io.FileIO.__init__(self, path, *args, **kwargs)

    def write(self, *args, **kwargs):
        """
        writes to io
        :param args:
        :param kwargs:
        :return:
        """
        progressbar(self.tell(), self._total_size)
        return io.FileIO.write(self, *args, **kwargs)

    def read(self, size=-1):
        """
        read from io
        :param size:
        :return:
        """
        progressbar(self.tell(), self._total_size)
        return io.FileIO.read(self, size)


tarfile.TarFile.fileobject = ProgressFileObject


def extract(tmpdir, tar_file_object):
    """
    extracts a tar file to the specified directory, with progressbar.
    :param tmpdir:
    :param tar_file_object:
    :return:
    """
    with tarfile.open(fileobj=ProgressFileObject(tar_file_object.name), mode='r') as tar:
        try:
            tar.extractall(path=tmpdir)
            print("")
            global global_char_pos
            global_char_pos = 0
            printc("Tar extraction completed", BColors.header)
        except KeyError:
            printc("tar file doesnt have the correct subdirs", BColors.fail)
        except Exception as e:
            printc("something went very wrong {}".format(str(e)), BColors.fail)


def set_hostname(tmpdir, hostname):
    """
    sets the hostname in /etc/hostname and mangles /etc/hosts to match
    :param tmpdir:
    :param hostname: hostname to set to.
    :return:
    """
    printc("Fixing hostname and hosts file:", BColors.blue)
    try:
        with open(os.path.join(tmpdir, "root", "etc", "hostname"), 'w') as f:
            f.write(hostname + "\n")

        with open(os.path.join(tmpdir, "root", "etc", "hosts"), 'w') as hosts_file:
            h_tmpl = "127.0.0.1\tlocalhost.localdomain localhost {hostname}\n"
            h_tmpl += "::1\tlocalhost.localdomain localhost {hostname}\n"
            hosts_file.write(h_tmpl.format(hostname=hostname))
    except Exception as e:
        printc("Couldnt fix hostname {}".format(str(e)), BColors.fail)
    printc(hostname, BColors.bold, BColors.blue)


def remove_torfiles(tmpdir):
    """
    removes files relevant to tor.
    This stops multiple raspberry pis having the same onion address.
    :param tmpdir:
    :return:
    """
    for path in glob.glob(os.path.join(tmpdir, "root", "home", "tor_private", "*")):
        os.remove(path)


def reset_machineid(tmpdir):
    """
    resets the machine-id on the raspberry pi.
    unfortunately this will only work on a machine with systemd.
    :param tmpdir:
    :return:
    """
    printc("Resetting machine-id", BColors.warn)
    machine_id_path = os.path.join(tmpdir, "root", "etc", "machine-id")
    try:
        os.remove(machine_id_path)
    except:
        pass
    cmd = "systemd-machine-id-setup --root {}".format(os.path.join(tmpdir, "root"))
    try:
        v = subprocess.check_output([cmd], shell=True, universal_newlines=True)
        printc(v, BColors.blue)
        with open(machine_id_path, 'r') as f:
            printc("machine-id: {}".format(f.read()), BColors.blue)
        printc("machine-id reset", BColors.header)
    except subprocess.CalledProcessError as e:
        printc("Couldnt call systemd-machine-id-setup: {}".format(str(e)), BColors.fail)
    except Exception as e:
        printc("Exception while resetting machine-id: {}".format(str(e)), BColors.fail)


def remove_ssh_keys(tmpdir):
    """
    removes ssh keys from the sd card.
    This stops new pis attempting to use ssh keys from another pi.
    :param tmpdir:
    :return:
    """
    printc("Removing old ssh keys", BColors.blue)
    ssh_files = glob.glob(os.path.join(tmpdir, "root", "home", ".ssh", "*"))
    for path in ssh_files:
        if not os.path.basename(path) == "key_token":
            os.remove(path)
        else:
            printc(".ssh dir has key token", BColors.green)


def remove_configs(tmpdir):
    """
    removes old configuration files.
    this stops raspberry pis reporting cameras that they do not have (the reporting is done per configuration file).
    :param tmpdir:
    :return:
    """
    printc("Removing old configs", BColors.blue)
    configs = glob.glob(os.path.join(tmpdir, "root", "home", "spc-eyepi", "configs_byserial", "*"))
    for path in configs:
        os.remove(path)


def copy_boot(tmpdir):
    """
    copys the files for the raspberry pi boot partition over
    :param tmpdir:
    :return:
    """
    try:
        boot = os.path.join(tmpdir, "boot")
        files = glob.glob(os.path.join(boot, "*"))
        newfiles = glob.glob(os.path.join(command_line_args.update_boot, "*"))
        for file in files:
            if os.path.isfile(file):
                os.remove(file)
            if os.path.isdir(file):
                shutil.rmtree(file)

        for file in newfiles:
            if os.path.isfile(file):
                shutil.copy(file, os.path.join(tmpdir, file))
            if os.path.isdir(file):
                shutil.copytree(file, os.path.join(tmpdir, file))
    except Exception as e:
        printc("Failed copying the boot dir: {}".format(str(e)))

def create_tarfile(tmpdir, fn):
    """
    creates a new tarfile from which to create new raspberry pis sd cards.
    :param tmpdir:
    :param fn:
    :return:
    """
    if os.path.isfile(fn):
        os.remove(fn)
    with open(fn, "w") as f:
        f.write("")
    with tarfile.open(mode="w", fileobj=ProgressFileObject(fn, mode='w')) as tar:
        global global_char_pos
        global_char_pos = 0
        printc("Adding /boot", BColors.header)
        tar.add(os.path.join(tmpdir, "boot"), "boot")
        global_char_pos = 0
        printc("Adding /", BColors.header)
        tar.add(os.path.join(tmpdir, "root"), "root")


def mkdirs_and_mount(tmpdir):
    """
    creates the appropriate directories and mounts the sd card partitions to them.
    creates a directory "boot" to mount the fat32 boot directory for pis ("/boot")
    creates a directory "root" to mount the ext4 directory to contain the OS ("/")

    :param tmpdir:
    :return bool: True if success, False if failed and shouldn't continue.
    """
    d = glob.glob(command_line_args.blockdevice[0] + "*")
    d.sort()
    d.pop(0)
    os.makedirs(os.path.join(tmpdir, "boot"), exist_ok=True)
    os.makedirs(os.path.join(tmpdir, "root"), exist_ok=True)
    if not len(d) >= 2:
        printc("Didn't find 2 partitions on the block device to mount...", BColors.fail)
        return False
    try:
        printc("Mounting {} to {}".format(" ".join(d), tmpdir), BColors.under)

        cmdstring = "mount {block_device} {path}"
        v = subprocess.run([cmdstring.format(block_device=d[0], path=os.path.join(tmpdir, "boot"))],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                           shell=True, universal_newlines=True)
        v.check_returncode()

        z = subprocess.run([cmdstring.format(block_device=d[1], path=os.path.join(tmpdir, "root"))],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                           shell=True, universal_newlines=True)
        z.check_returncode()
        printc("Mounted", BColors.under)
        return True
    except subprocess.CalledProcessError as e:
        printc("Couldn't call mount {}\n{}".format(str(e), str(e.output)), BColors.fail)
    except Exception as e:
        printc("Something went wrong mounting: {}".format(str(e)), BColors.fail)
    sync_unmount()
    return False


def fix_boot(tmpdir):
    with open(os.path.join(tmpdir,"boot","config.txt"), 'w') as f:
        f.write("hdmi_force_hotplug=1\ngpu_mem=256\nconsoleblank=0\ndisable_camera_led=1\nstart_x=1")


def sync_unmount():
    """
    unmounts block device and syncs filesystem
    :return:
    """
    printc("Syncing fs", BColors.under)
    os.sync()

    printc("Unmounting {}".format(command_line_args.blockdevice[0]), BColors.under)
    command_string = "umount {}?"
    try:
        subprocess.check_output([command_string.format(command_line_args.blockdevice[0])],
                                stderr=subprocess.STDOUT,
                                shell=True, universal_newlines=True)
        printc("Unmounted {}".format(command_line_args.blockdevice[0]), BColors.under)
    except subprocess.CalledProcessError as e:
        printc("Couldnt call umount: {}\n{}".format(str(e), str(e.output)), BColors.fail)
    except Exception as e:
        printc("Some other exception occured during the unmounting process: {}".format(str(e)), BColors.fail)


if __name__ == '__main__':

    if command_line_args.tarfile:
        printc("Formatting", BColors.warn)
        partition_and_format()

    with TemporaryDirectory() as temp_dir:
        try:
            if not mkdirs_and_mount(temp_dir):
                printc("Couldnt mount and create directories, quitting.")
                sys.exit(1)

            if command_line_args.tarfile:
                printc("Extracting {}".format(command_line_args.tarfile.name), BColors.header)
                extract(temp_dir, command_line_args.tarfile)
                update_via_github(temp_dir)
                set_hostname(temp_dir, gname)
                reset_machineid(temp_dir)
                remove_torfiles(temp_dir)
                remove_ssh_keys(temp_dir)
                remove_configs(temp_dir)

            if command_line_args.update_boot:
                printc("Updating boot partition with whatever is in {}".format(command_line_args.update_boot), BColors.header)
                copy_boot(temp_dir)

            if command_line_args.to_tarfile:
                printc("Creating new tarfile", BColors.header)
                update_via_github(temp_dir)
                set_hostname(temp_dir, "BLANK_FLASH")
                remove_torfiles(temp_dir)
                remove_ssh_keys(temp_dir)
                remove_configs(temp_dir)
                create_tarfile(temp_dir, command_line_args.to_tarfile)

            if command_line_args.update:
                printc("Updating", BColors.blue)
                update_via_github(temp_dir)
                set_hostname(temp_dir, gname)


            if command_line_args.restore:
                printc("Restoring", BColors.blue)
                restore(temp_dir, bakdir=command_line_args.restore)

            elif command_line_args.backup:
                printc("Backing up", BColors.blue)
                backup(temp_dir)

            if command_line_args.api_token:
                printc("Writing api token", BColors.blue)
                write_key_token(temp_dir)

            if os.path.isfile("db"):
                printc("Writing db file", BColors.blue)
                shutil.copy("db", os.path.join(temp_dir, "root", "home", "spc-eyepi", "db"))

            if os.path.isfile("db_private"):
                printc("Writing private db file", BColors.blue)
                shutil.copy("db_private", os.path.join(temp_dir, "root", "home", "spc-eyepi", "db"))

            if command_line_args.remove_tor:
                remove_torfiles(temp_dir)

            if command_line_args.remove_configs:
                remove_configs(temp_dir)

            if command_line_args.reset_machine_id:
                reset_machineid(temp_dir)

            if command_line_args.remove_keys:
                remove_ssh_keys(temp_dir)

            fix_boot(temp_dir)
        except Exception as e:
            printc("Unhandled exception: {}".format(str(e)), BColors.fail)
        finally:
            sync_unmount()