Modifies a raspbian image file to set static networking parameters and a hostname, as well as turning on SSH.

modify_raspbian_image.py

# setuppi.py
# Author: twitchyliquid64
# Automatically modifies a raspbian .img file to:
# - Set a static IP address, gateway, and DNS server
# - Enable SSH
# - Set a hostname.
#
# Tested and working with a default install of Ubuntu 16.04
# USAGE: python setuppi.py <imagefile>
# EG:    sudo python setuppi.py 2017-04-10-raspbian-jessie-lite.img
# The program will prompt for information.

import os, sys, pwd, time
from subprocess import call, check_output
import shutil

DHCPD_CONF = '''
# A sample configuration for dhcpcd.
# See dhcpcd.conf(5) for details.

# Allow users of this group to interact with dhcpcd via the control socket.
#controlgroup wheel

# Inform the DHCP server of our hostname for DDNS.
hostname

# Use the hardware address of the interface for the Client ID.
clientid
# or
# Use the same DUID + IAID as set in DHCPv6 for DHCPv4 ClientID as per RFC4361.
#duid

# Persist interface configuration when dhcpcd exits.
persistent

# Rapid commit support.
# Safe to enable by default because it requires the equivalent option set
# on the server to actually work.
option rapid_commit

# A list of options to request from the DHCP server.
option domain_name_servers, domain_name, domain_search, host_name
option classless_static_routes
# Most distributions have NTP support.
option ntp_servers
# Respect the network MTU.
# Some interface drivers reset when changing the MTU so disabled by default.
#option interface_mtu

# A ServerID is required by RFC2131.
require dhcp_server_identifier
./2017-04-10-raspbian-jessie-lite.img
nohook lookup-hostname

'''

class bcolors: #Thanks stackoverflow
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def runCommand(args):
    print(bcolors.OKGREEN + ' '.join(args) + bcolors.ENDC)
    call(args)

# returns array of [sectorStart, sectionName, Type]
def fdisk(fname):
    cmd = ['fdisk', '-l', '-o', 'start,type,device', fname]
    print(bcolors.OKGREEN + ' '.join(cmd) + bcolors.ENDC)
    output = check_output(cmd).split('\n\n')[1]
    partitions = []
    for line in output.split('\n')[1:]:
        line_tokens = line.strip().split(' ')
        if line_tokens[0] == '':
            continue
        out = [line_tokens[0]]

        line = ''
        for token in line_tokens[1:]:
            if not token.startswith(fname):
                line += token
            else:
                out.append(token)
        out.append(line)
        partitions.append(out)
    return partitions

if __name__ == "__main__":
    if os.geteuid() != 0:
        print("Need root! Please run with sudo.")
        sys.exit(1)

    if len(sys.argv) < 2:
        print("USAGE: %s <raspbian-image>" % sys.argv[0])
        sys.exit(1)

    ip = raw_input("Please enter the IP/subnet combination (EG: 192.168.0.15/24): ")
    gateway = raw_input("Please enter the gateway IP (EG: 192.168.0.1): ")
    dns = raw_input("Please enter the IP of the DNS server (EG: 8.8.8.8): ")
    hostname = raw_input("Please enter the hostname that the Pi should be (EG: mybox): ")

    try:
        shutil.rmtree('fat_mnt')
        shutil.rmtree('ext4_mnt')
    except:
        pass

    os.mkdir('fat_mnt')
    os.mkdir('ext4_mnt')

    partitions = fdisk(sys.argv[1])
    for partition in partitions:
        if partition[2] == 'W95FAT32(LBA)':
            print 'Mounting FAT32 partition starting at sector %s to ./%s' % (partition[0], 'fat_mnt')
            runCommand(['mount', '-t', 'vfat', '-o', 'loop,offset=' + str(512*int(partition[0])), sys.argv[1], 'fat_mnt'])
        if partition[2] == 'Linux':
            print 'Mounting EXT4 partition starting at sector %s to ./%s' % (partition[0], 'ext4_mnt')
            runCommand(['mount', '-o', 'loop,offset=' + str(512*int(partition[0])), sys.argv[1], 'ext4_mnt'])

    #with open('ext4_mnt/etc/passwd', 'rw') as user_file:
    #    print(user_file.read())

    print(bcolors.OKBLUE + "Writing /etc/dhcpcd.conf:" + bcolors.ENDC)
    print(bcolors.OKBLUE + "\tIP + subnet: " + ip)
    print(bcolors.OKBLUE + "\tGateway:" + gateway + bcolors.ENDC)
    print(bcolors.OKBLUE + "\tDNS: " + dns + bcolors.ENDC)
    with open('ext4_mnt/etc/dhcpcd.conf', 'w') as dhcpd_file:
        dhcpd_file.write(DHCPD_CONF)
        dhcpd_file.write('\ninterface eth0\n\n')
        dhcpd_file.write('static ip_address=%s\n' % ip)
        dhcpd_file.write('static routers=%s\n' % gateway)
        dhcpd_file.write('static domain_name_servers=%s\n' % dns)

    print(bcolors.OKBLUE + "Enabling SSH" + bcolors.ENDC)
    runCommand(['touch', 'fat_mnt/ssh'])

    print(bcolors.OKBLUE + "Setting hostname to '%s'" % hostname + bcolors.ENDC)
    with open('ext4_mnt/etc/hostname', 'w') as hostname_file:
        hostname_file.write(hostname)

    time.sleep(2)
    runCommand(['umount', 'fat_mnt'])
    runCommand(['umount', 'ext4_mnt'])

    try:
        shutil.rmtree('fat_mnt')
        shutil.rmtree('ext4_mnt')
    except:
        pass

    print("Image modifications complete. Flash to your Pi SD card like: " + bcolors.UNDERLINE + "sudo dd bs=1M if=<imagefile> of=/dev/sdd" + bcolors.ENDC)
    print("GOTCHA: After writing the changes, eject via your file manager AND WAIT for it to explicitly say it is done writing files. Otherwise you corrupt shit.")