001_setup_sd_card

#!/usr/bin/env python3

import argparse
import os
import sys
import subprocess
import pathlib
import time

BASH_EXECUTABLE = '/bin/bash'

def parseArgs():
    parser = argparse.ArgumentParser(
        description=("Setup new sd card for raspberry pi. Must be run as sudo. Images disk, sets up wifi, sets up " +
            "ssh. Run from laptop."),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('--disk-num', dest='disk_num', action='store', required = True,
        help='Disk number via `diskutil list`. This disk will be overwritten, so don\'t get this wrong.')
    parser.add_argument('--disk-image', dest='disk_image', action='store', required = True,
        help='Path to disk image. Download Raspberry Pi OS Lite image from https://www.raspberrypi.org/software/operating-systems/')
    parser.add_argument('--device-path', dest='device_path', action='store', default='/Volumes/bootfs',
        help='Mounted volume path after SD card has been imaged.')
    parser.add_argument('--user-name', dest='user_name', action='store', required = True,
        help='User account to create')
    parser.add_argument('--user-password', dest='user_password', action='store', required = True,
        help='Password for the user account')
    parser.add_argument('--wifi-network-name', dest='wifi_network_name', action='store', required = True,
        help='Name of the network to join')
    parser.add_argument('--wifi-password', dest='wifi_password', action='store', required = True,
        help='Password of the network to join')
    parser.add_argument('--wifi-country-code', dest='wifi_country_code', action='store', default='US',
        help='Your ISO-3166-1 two-letter country code. See: ' +
        'https://www.iso.org/obp/ui/#search')

    if os.getuid() != 0:
        print("This script must be run as sudo.\n", file = sys.stderr)
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    args.device_path = args.device_path.rstrip('/')
    return args

def validateDependencies():
    try:
        import passlib
    except ImportError:
        print("The passlib package is required. Please install it: " +
            "'sudo -H PIP_BREAK_SYSTEM_PACKAGES=1 python3 -m pip install --upgrade passlib'.", file = sys.stderr)
        sys.exit(1)

# https://www.raspberrypi.org/documentation/installation/installing-images/mac.md
def imageSdCard():
    print("Writing disk image, may take up to 15 minutes...")
    command(f'diskutil unmountDisk /dev/disk{args.disk_num}')
    command(f'sudo dd bs=1m if={args.disk_image} of=/dev/rdisk{args.disk_num}; sync')

    time.sleep(5)

    # Wait a few seconds for args.device_path to remount again (I guess this takes some time?).
    # Need it to be remounted before writing wifi and ssh files.
    time_limit = time.time() + 10
    while True:
        if os.path.isdir(args.device_path):
            break
        time.sleep(0.5)
        if time.time() > time_limit:
            print(f"Failed to remount {args.device_path}.", file = sys.stderr)
            sys.exit(1)

    # prevent weirdness where subsequent writes are flakey?
    time.sleep(5)

# https://www.raspberrypi.com/documentation/computers/configuration.html#configuring-a-user
# https://www.raspberrypi.com/news/raspberry-pi-bullseye-update-april-2022/
def setupUserAccount():
    print("Setting up user account...")
    import passlib.hash

    # Use fewer hashing rounds than the default of 656000. This makes password based login faster. In testing,
    # it reduced the time to log in on a pi4 using a password from ~8s to ~1s. Furthermore, this is the number
    # of rounds that would be used by default if following the raspberry pi official docs for generating
    # passwords.
    #
    # See: https://passlib.readthedocs.io/en/stable/lib/passlib.hash.sha512_crypt.html#passlib.hash.sha512_crypt
    password_hash = passlib.hash.sha512_crypt.using(rounds=5000).hash(args.user_password)
    f = None
    try:
        f = open(f"{args.device_path}/userconf.txt", "w")
    except PermissionError:
        print('This script must be run as sudo', file = sys.stderr)
        sys.exit(1)

    f.write(f'{args.user_name}:{password_hash}')
    f.close()

# https://raspberrypi.stackexchange.com/a/57023
def setupWifi():
    print("Setting up wifi...")
    my_dir = pathlib.Path(__file__).parent.absolute()
    command(f'sudo {my_dir}/wifi_config --wifi-network-name {args.wifi_network_name} ' +
        f'--wifi-password {args.wifi_password} --wifi-country-code {args.wifi_country_code} ' +
        f'--device-path {args.device_path}')

def setupSsh():
    print("Enabling ssh...")
    command(f'sudo touch {args.device_path}/ssh')

def validate():
    try:
        command(f'sudo ls -l {args.device_path}/cmdline.txt')
    except Exception:
        print("ERROR: failed to validate raspbian install", file = sys.stderr)
        sys.exit(1)

    try:
        command(f'sudo ls -l {args.device_path}/ssh')
    except Exception:
        print("ERROR: failed to validate ssh setup", file = sys.stderr)
        sys.exit(1)

    try:
        command(f'sudo ls -l {args.device_path}/wpa_supplicant.conf')
    except Exception:
        print("ERROR: failed to validate wifi setup", file = sys.stderr)
        sys.exit(1)

def finish():
    print("Ejecting SD card...")
    command(f'sudo diskutil eject /dev/rdisk{args.disk_num}')
    print(f"""
SD card setup was successful! Please insert the SD card into the raspberry pi and boot it up.
Your raspberry pi will be accessible via:

    ssh {args.user_name}@raspberrypi.local
    password: {args.user_password}
""")

def command(command):
    return (subprocess
        .check_output(command, shell = True, executable = BASH_EXECUTABLE)
        .decode("utf-8"))


args = parseArgs()
validateDependencies()
imageSdCard()
setupUserAccount()
setupWifi()
setupSsh()
validate()
finish()