Commit 5d489031 authored by segfault's avatar segfault
Browse files

Fix and improve tails-get-bootinfo.py

 - Fix doctest
 - Fix passing end=None to print()
 - Use open() instead of sh.cat()
 - Use /usr/bin/env in shebang
 - Make docstrings PEP 257 compliant
 - Other style improvements
parent 8910940d
#! /usr/bin/python3
'''
Get Tails boot info.
#! /usr/bin/env python3
"""
Get Tails boot info.
Test with "python3 tails-get-bootinfo.py doctest".
Test with "python3 tails-get-bootinfo.py doctest".
goodcrypto.com converted from bash to python and added basic tests.
>>> import sh
>>> sh.Command(sys.argv[0])('kernel')
/lib/live/mount/medium/live/vmlinuz
>>> sh.Command(sys.argv[0])('initrd')
/lib/live/mount/medium/live/initrd.img
>>> sh.Command(sys.argv[0])(_ok_code=(1))
Usage: tails-get-bootinfo kernel|initrd
<BLANKLINE>
"""
goodcrypto.com converted from bash to python and added basic tests.
'''
import os
import sys
import sh
LIVE_IMAGE_MOUNTPOINT = '/lib/live/mount/medium'
SCRIPT_NAME = os.path.basename(sys.argv[0])
def main(args):
kernel = initrd = None
result = sh.cat('/proc/cmdline').stdout.decode()
for line in result.split(' '):
if line.startswith('BOOT_IMAGE='):
kernel = line[len('BOOT_IMAGE='):]
elif line.startswith('initrd='):
initrd = line[len('initrd='):]
def main(*args):
kernel = None
initrd = None
with open('/proc/cmdline') as f:
kernel_params = f.read().split()
for param in kernel_params:
if param.startswith('BOOT_IMAGE='):
kernel = param[len('BOOT_IMAGE='):]
elif param.startswith('initrd='):
initrd = param[len('initrd='):]
# Sanity checks
if not kernel:
sys.exit(4)
if not initrd:
sys.exit(5)
if not kernel or not initrd:
print('Failed to parse /proc/cmdline', file=sys.stderr)
if args[1] == 'kernel':
print(LIVE_IMAGE_MOUNTPOINT + kernel, end=None)
elif args[1] == 'initrd':
print(LIVE_IMAGE_MOUNTPOINT + initrd, end=None)
if 'kernel' in args:
print(LIVE_IMAGE_MOUNTPOINT + kernel, end="")
elif 'initrd' in args:
print(LIVE_IMAGE_MOUNTPOINT + initrd, end="")
else:
print('Usage: {} kernel|initrd PATH'.format(SCRIPT_NAME), file=sys.stderr)
sys.exit(3)
print('Usage: tails-get-bootinfo kernel|initrd PATH', file=sys.stderr)
sys.exit(1)
'''
>>> def test_boot(boot):
... """ Run this script. """
...
... common_args = ['BOOT_IMAGE=/kernel-boot-path', 'initrd=/initrd-boot-path']
... command_path = os.path.join(os.getcwd(), sys.argv[0])
... this_command = sh.Command(command_path)
... out = this_command(boot, *common_args).stdout.decode().strip()
... print(out)
>>> test_boot('kernel')
b'/lib/live/mount/medium/kernel-boot-path'
>>> test_boot('initrd')
b'/lib/live/mount/medium/initrd-boot-path'
'''
if __name__ == "__main__":
if sys.argv and len(sys.argv) > 1:
if sys.argv[1] == 'doctest':
from doctest import testmod
testmod()
else:
main(sys.argv)
sys.exit(0)
if len(sys.argv) > 1 and sys.argv[1] == 'doctest':
import doctest
doctest.testmod()
elif len(sys.argv) == 2:
main(*sys.argv[1:])
else:
print('Usage: tails-get-bootinfo kernel|initrd')
sys.exit(-1)
sys.exit(1)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment