|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import os |
| 4 | +import re |
| 5 | +import subprocess |
| 6 | +import time |
| 7 | +import sys |
| 8 | +from tabulate import tabulate |
| 9 | + |
| 10 | +HELP = "Usage:\n\ |
| 11 | +\tpxetools --add serial\n\ |
| 12 | +\tpxetools --remove serial\n\ |
| 13 | +\tpxetools --list\n" |
| 14 | + |
| 15 | +LAN = None |
| 16 | +NFS_IP = None |
| 17 | + |
| 18 | +def cmd(cmd, print_out=False, print_cmd=True, can_fail=False): |
| 19 | + if print_cmd: |
| 20 | + print("# {}".format(cmd)) |
| 21 | + |
| 22 | + out = None |
| 23 | + |
| 24 | + try: |
| 25 | + out = subprocess.check_output(cmd, shell=True).decode() |
| 26 | + except: |
| 27 | + if not can_fail: |
| 28 | + raise |
| 29 | + |
| 30 | + if print_out: |
| 31 | + print(out) |
| 32 | + |
| 33 | + return out |
| 34 | + |
| 35 | +def add(): |
| 36 | + serial = None |
| 37 | + if len(sys.argv) != 3: |
| 38 | + if os.path.exists("/usr/local/sbin/get_serial"): |
| 39 | + # Auto detect |
| 40 | + print("Will try and auto detect serial. Please plug in your Pi now!") |
| 41 | + cmd("systemctl stop dnsmasq") |
| 42 | + cmd("iptables -t raw --flush") |
| 43 | + |
| 44 | + try: |
| 45 | + serial = cmd("get_serial", print_cmd=False).rstrip() |
| 46 | + except KeyboardInterrupt: |
| 47 | + cmd("iptables-restore < /etc/iptables/rules.v4") |
| 48 | + cmd("systemctl start dnsmasq") |
| 49 | + sys.exit() |
| 50 | + |
| 51 | + print("Found serial") |
| 52 | + cmd("iptables-restore < /etc/iptables/rules.v4") |
| 53 | + cmd("systemctl start dnsmasq") |
| 54 | + else: |
| 55 | + serial = sys.argv[2] |
| 56 | + |
| 57 | + # Validate serial |
| 58 | + serial = serial.lstrip("0") |
| 59 | + if not re.search("^[0-9a-f]{8}$", serial): |
| 60 | + raise Exception("Invalid serial number {}".format(serial)) |
| 61 | + |
| 62 | + print("Serial: {}".format(serial)) |
| 63 | + owner = input("Owner Name (ex Gordon): ") |
| 64 | + name = input("Name for pi: ") |
| 65 | + |
| 66 | + print("Select a base image:") |
| 67 | + selection = ["I will prepare my own filesystem"] + os.listdir('/nfs/bases') |
| 68 | + for i in range (0, len(selection)): |
| 69 | + print("\t{}. {}".format(i + 1, selection[i])) |
| 70 | + |
| 71 | + img_choice = input("Enter an option number: ") |
| 72 | + |
| 73 | + tftp_path = "/tftpboot/{}".format(serial) |
| 74 | + nfs_path = "/nfs/{}".format(serial) |
| 75 | + if os.path.exists(tftp_path): |
| 76 | + raise Exception("{} already exists!".format(tftp_path)) |
| 77 | + |
| 78 | + valid_img_choice = True |
| 79 | + try: |
| 80 | + img_choice = int(img_choice) |
| 81 | + # Make 0 based again |
| 82 | + img_choice -= 1 |
| 83 | + except ValueError: |
| 84 | + valid_img_choice = False |
| 85 | + |
| 86 | + if img_choice < 0 or img_choice > (len(selection) - 1) or (not valid_img_choice): |
| 87 | + raise Exception("Invalid image choice {}".format(img_choice)) |
| 88 | + |
| 89 | + img = None |
| 90 | + if img_choice > 0: |
| 91 | + img = "/nfs/bases/{}".format(selection[img_choice]) |
| 92 | + |
| 93 | + print("\nSetting up pi...\n") |
| 94 | + |
| 95 | + cmd("iptables -t raw -I DHCP_clients -m mac --mac-source {0} -j ACCEPT".format(mac(serial))) |
| 96 | + cmd("iptables-save > /etc/iptables/rules.v4") |
| 97 | + cmd("mkdir {}".format(tftp_path)) |
| 98 | + cmd("cp -r /tftpboot/base/* {}".format(tftp_path)) |
| 99 | + cmd("mkdir {}".format(nfs_path)) |
| 100 | + cmd("echo \"{} *(rw,sync,no_subtree_check,no_root_squash)\" >> /etc/exports".format(nfs_path)) |
| 101 | + cmd("exportfs -a") |
| 102 | + |
| 103 | + cmdline_txt = "dwc_otg.lpm_enable=0 root=/dev/nfs nfsroot={}:{} rw ip=dhcp rootwait elevator=deadline".format(NFS_IP, nfs_path) |
| 104 | + cmd("echo \"{}\" > {}/cmdline.txt".format(cmdline_txt, tftp_path)) |
| 105 | + cmd("echo \"{}\" > {}/owner".format(owner, tftp_path)) |
| 106 | + cmd("echo \"{}\" > {}/name".format(name, tftp_path)) |
| 107 | + |
| 108 | + if img_choice: |
| 109 | + cmd("mkdir -p /mnt/tmp") |
| 110 | + cmd("kpartx -a -v {}".format(img), print_out=True) |
| 111 | + time.sleep(1) |
| 112 | + cmd("mount /dev/mapper/loop0p2 /mnt/tmp") |
| 113 | + cmd("rsync -a /mnt/tmp/ {}/".format(nfs_path)) |
| 114 | + cmd("umount /mnt/tmp") |
| 115 | + cmd("echo > {}/etc/fstab".format(nfs_path)) |
| 116 | + cmd("cd {}/etc/init.d; rm dhcpcd dphys-swapfile raspi-config resize2fs_once".format(nfs_path)) |
| 117 | + cmd("cd {}/etc/systemd/system; rm -r dhcp* multi-user.target.wants/dhcp*".format(nfs_path)) |
| 118 | + cmd("mount /dev/mapper/loop0p1 /mnt/tmp") |
| 119 | + cmd("rsync -a --exclude bootcode.bin --exclude start.elf --exclude cmdline.txt /mnt/tmp/ {}/".format(tftp_path)) |
| 120 | + cmd("umount /mnt/tmp") |
| 121 | + cmd("kpartx -d -v {}".format(img), print_out=True) |
| 122 | + else: |
| 123 | + print("You opted to prep your own system so please put files in:\n\t{}\n\t{}".format(tftp_path, nfs_path)) |
| 124 | + print("I have wrote you a cmdline.txt so don't change it:") |
| 125 | + print(cmdline_txt) |
| 126 | + |
| 127 | + print("Should be working. You might have to wait a couple of minutes before SSH lets you in") |
| 128 | + |
| 129 | +def remove(): |
| 130 | + serial = sys.argv[2] |
| 131 | + |
| 132 | + if not re.search("^[0-9a-f]{8}$", serial): |
| 133 | + raise Exception("Invalid serial number {}".format(serial)) |
| 134 | + |
| 135 | + sure = "" |
| 136 | + while sure != "Y" and sure != "N": |
| 137 | + sure = input("ARE YOU SURE YOU WANT TO DELETE {}? Y/N: ".format(serial)) |
| 138 | + |
| 139 | + if sure == "N": |
| 140 | + print("Aborting") |
| 141 | + return |
| 142 | + |
| 143 | + cmd("rm -rf /tftpboot/{}".format(serial), can_fail=True) |
| 144 | + |
| 145 | + nfspath = "/nfs/{}".format(serial) |
| 146 | + cmd("rm -rf {}".format(nfspath), can_fail=True) |
| 147 | + |
| 148 | + print("Removing from /etc/exports") |
| 149 | + with open("/etc/exports", 'r+') as f: |
| 150 | + export = f.read() |
| 151 | + f.seek(0) |
| 152 | + |
| 153 | + for l in export.splitlines(): |
| 154 | + if l.startswith(nfspath) == False: |
| 155 | + f.write(l + "\n") |
| 156 | + |
| 157 | + f.truncate() |
| 158 | + |
| 159 | + mac_uppercase = mac(serial).upper() |
| 160 | + print("Removing from iptables rules") |
| 161 | + with open("/etc/iptables/rules.v4", 'r+') as f: |
| 162 | + rules = f.read() |
| 163 | + f.seek(0) |
| 164 | + |
| 165 | + for l in rules.splitlines(): |
| 166 | + if mac_uppercase not in l: |
| 167 | + f.write(l + "\n") |
| 168 | + |
| 169 | + f.truncate() |
| 170 | + |
| 171 | + cmd("iptables-restore < /etc/iptables/rules.v4") |
| 172 | + |
| 173 | +def mac(serial): |
| 174 | + # MAC is least significant bits of serial |
| 175 | + # Example: |
| 176 | + # b8:27:eb:be:e3:d2 |
| 177 | + # 0000000044bee3d2 |
| 178 | + prefix = "b8:27:eb:" |
| 179 | + suffix = serial[-6:] |
| 180 | + |
| 181 | + # Insert colon every 2 chars. Stolen from stackoverflow |
| 182 | + s = suffix |
| 183 | + s = ":".join(a+b for a,b in zip(s[::2], s[1::2])) |
| 184 | + return prefix + s |
| 185 | + |
| 186 | +def ip(mac): |
| 187 | + # Hacky way of getting IP from Mac Address |
| 188 | + cmdstr = ("nmap -sn -PR --min-rate 5000 --max-retries=0 {0} " |
| 189 | + "| grep -i -B 2 {1} | head -n 1 |" |
| 190 | + " awk '{{print $5}}'").format(LAN, mac) |
| 191 | + ip = cmd(cmdstr, print_cmd=False) |
| 192 | + if len(ip) == 0: |
| 193 | + ip = "Not found" |
| 194 | + else: |
| 195 | + ip = ip.rstrip() |
| 196 | + |
| 197 | + return ip |
| 198 | + |
| 199 | +def file_or_none(base, file): |
| 200 | + path = os.path.join(base, file) |
| 201 | + if not os.path.exists(path): |
| 202 | + return "Not Found" |
| 203 | + |
| 204 | + with open(path, 'r') as fh: |
| 205 | + return fh.read() |
| 206 | + |
| 207 | +def list(): |
| 208 | + tbl = [["Serial", "Owner", "Name", "MAC", "IP"]] |
| 209 | + pis = os.listdir("/tftpboot") |
| 210 | + pis = [p for p in pis if p != "base" and p != "bootcode.bin"] |
| 211 | + |
| 212 | + for p in pis: |
| 213 | + base = os.path.join("/tftpboot", p) |
| 214 | + tbl.append(["0x{}".format(p), file_or_none(base, "owner"), file_or_none(base, "name"), mac(p), ip(mac(p))]) |
| 215 | + |
| 216 | + print(tabulate(tbl, headers="firstrow")) |
| 217 | + |
| 218 | +if __name__ == "__main__": |
| 219 | + if os.geteuid() != 0: |
| 220 | + exit("Please run with sudo or as root user") |
| 221 | + |
| 222 | + LAN = cmd('ip -4 addr show dev eth0 | grep inet | cut -d " " -f6 | xargs ipcalc | grep Network | cut -d " " -f4', print_cmd=False).rstrip() |
| 223 | + NFS_IP = cmd('ifconfig eth0 | grep "inet addr" | cut -d " " -f 12 | cut -d ":" -f 2', print_cmd=False).rstrip() |
| 224 | + |
| 225 | + if len(sys.argv) == 1: |
| 226 | + print(HELP) |
| 227 | + elif sys.argv[1] == "--list": |
| 228 | + list() |
| 229 | + elif sys.argv[1] == "--add": |
| 230 | + add() |
| 231 | + elif sys.argv[1] == "--remove": |
| 232 | + remove() |
| 233 | + else: |
| 234 | + print(HELP) |
0 commit comments