stage3img: Remove (superseded by gentoo-img)
parent
2c5d930d0b
commit
261ea1fe16
598
stage3img.py
598
stage3img.py
|
|
@ -1,598 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import collections
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import string
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
|
||||
log = logging.getLogger('stage3img')
|
||||
|
||||
|
||||
LVM_PV_PART = {
|
||||
'type': 0x8e00,
|
||||
}
|
||||
|
||||
UNITS = {
|
||||
'k': 2 ** 10,
|
||||
'kb': 10 ** 3,
|
||||
'kib': 2 ** 10,
|
||||
'm': 2 ** 20,
|
||||
'mb': 10 ** 6,
|
||||
'mib': 2 ** 20,
|
||||
'g': 2 ** 30,
|
||||
'gb': 10 ** 9,
|
||||
'gib': 2 ** 30,
|
||||
't': 2 ** 40,
|
||||
'tb': 10 ** 12,
|
||||
'tib': 2 ** 40,
|
||||
'p': 2 ** 50,
|
||||
'pb': 10 ** 15,
|
||||
'pib': 2 ** 50,
|
||||
'e': 2 ** 60,
|
||||
'eb': 10 ** 18,
|
||||
'eib': 2 ** 60,
|
||||
'z': 2 ** 70,
|
||||
'zb': 10 ** 21,
|
||||
'zib': 2 ** 70,
|
||||
'y': 2 ** 80,
|
||||
'yb': 10 ** 24,
|
||||
'yib': 2 ** 80,
|
||||
}
|
||||
|
||||
SIZE_RE = re.compile(
|
||||
r'^(?P<value>[0-9]+(?:\.[0-9]+)?)\s*'
|
||||
r'(?P<unit>\w+)?\s*$'
|
||||
)
|
||||
|
||||
TRY_SSH_KEYS = [
|
||||
os.path.expanduser('~/.ssh/id_ed25519.pub'),
|
||||
os.path.expanduser('~/.ssh/id_rsa.pub'),
|
||||
os.path.expanduser('~/.ssh/id_ecdsa.pub'),
|
||||
]
|
||||
|
||||
|
||||
class CommandError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class MaxLevelFilter(logging.Filter):
|
||||
'''Filter log records above a particular level
|
||||
|
||||
:param max_level: The maximum level of log message to include
|
||||
'''
|
||||
|
||||
def __init__(self, max_level=logging.WARNING, *args, **kwargs):
|
||||
logging.Filter.__init__(self, *args, **kwargs)
|
||||
self.max_level = max_level
|
||||
|
||||
def filter(self, record):
|
||||
return record.levelno <= self.max_level
|
||||
|
||||
|
||||
class Image(object):
|
||||
|
||||
EXTRA_MOUNTS = [
|
||||
('/tmp', 'tmpfs', '-t', 'tmpfs'),
|
||||
('/run', 'tmpfs', '-t', 'tmpfs'),
|
||||
('/dev', '/dev', '--rbind'),
|
||||
('/proc', 'proc', '-t', 'proc'),
|
||||
('/sys', '/sys', '--rbind'),
|
||||
]
|
||||
SCRIPT_ENV = {
|
||||
'PATH': os.pathsep.join((
|
||||
'/usr/local/sbin',
|
||||
'/usr/sbin',
|
||||
'/sbin',
|
||||
'/usr/local/bin',
|
||||
'/usr/bin',
|
||||
'/bin',
|
||||
)),
|
||||
}
|
||||
|
||||
default_fstype = 'ext4'
|
||||
|
||||
def __init__(self, filename, fd=None):
|
||||
self.fd = fd
|
||||
self.filename = filename
|
||||
self.tempname = self.filename + '.tmp'
|
||||
self.loopdev = None
|
||||
self.partitions = None
|
||||
self.vgname = None
|
||||
self.volumes = None
|
||||
self.mountpoint = None
|
||||
self._extra_mounted = False
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, tb):
|
||||
self.unmount()
|
||||
self.deactivate_vg()
|
||||
self.detach_loopback()
|
||||
if self.fd is not None:
|
||||
os.close(self.fd)
|
||||
if exc_type:
|
||||
log.error('Error occurred, removing image')
|
||||
try:
|
||||
os.unlink(self.tempname)
|
||||
except:
|
||||
log.exception('Failed to remove temporary image file:')
|
||||
try:
|
||||
os.unlink(self.filename)
|
||||
except:
|
||||
log.exception('Failed to remove image file:')
|
||||
else:
|
||||
os.rename(self.tempname, self.filename)
|
||||
|
||||
|
||||
@property
|
||||
def filesystems(self):
|
||||
if self.partitions:
|
||||
partnum = 1
|
||||
for part in self.partitions:
|
||||
if 'mountpoint' not in part:
|
||||
continue
|
||||
dev = '{}p{}'.format(self.loopdev, partnum)
|
||||
yield (
|
||||
part['mountpoint'],
|
||||
dev,
|
||||
part.get('fstype'),
|
||||
part.get('fsopts', 'defaults'),
|
||||
)
|
||||
partnum += 1
|
||||
if self.volumes:
|
||||
for vol in self.volumes:
|
||||
if 'mountpoint' not in vol:
|
||||
continue
|
||||
dev = '/dev/{}/{}'.format(self.vgname, vol['name'])
|
||||
yield (
|
||||
vol['mountpoint'],
|
||||
dev,
|
||||
vol.get('fstype'),
|
||||
vol.get('fsopts', 'defaults'),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create(cls, filename):
|
||||
fd = os.open(filename, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
||||
return cls(filename, fd)
|
||||
|
||||
def resize(self, size):
|
||||
with open(self.tempname, 'wb') as f:
|
||||
f.truncate(size)
|
||||
|
||||
def attach_loopback(self):
|
||||
self.loopdev = run_cmd('losetup', '-f', '--show', self.tempname)
|
||||
log.info('Connected {} to {}'.format(self.tempname, self.loopdev))
|
||||
|
||||
def detach_loopback(self):
|
||||
if not self.loopdev:
|
||||
return
|
||||
log.info('Detaching {}'.format(self.loopdev))
|
||||
run_cmd('losetup', '-d', self.loopdev)
|
||||
|
||||
def partition(self, partitions):
|
||||
if not self.loopdev:
|
||||
self.attach_loopback()
|
||||
log.info('Partitioning {} with GPT'.format(self.loopdev))
|
||||
cmd = ['sgdisk', '-a', '4096', '-Z', '-g']
|
||||
for idx, part in enumerate(partitions):
|
||||
partnum = idx + 1
|
||||
cmd += ('-n', '{}::{}'.format(partnum, part.get('size', '')))
|
||||
if 'type' in part:
|
||||
cmd += ('-t', '{}:{:X}'.format(partnum, part['type']))
|
||||
if 'name' in part:
|
||||
cmd += ('-c', '{}:{}'.format(partnum, part['name']))
|
||||
cmd.append(self.loopdev)
|
||||
run_cmd_logged(*cmd)
|
||||
self.partitions = partititons
|
||||
|
||||
def setup_lvm(self, vgname, volumes):
|
||||
if not self.loopdev:
|
||||
self.attach_loopback()
|
||||
pvscan(self.loopdev)
|
||||
if self.partitions:
|
||||
pv = '{}p{}'.format(self.loopdev, len(self.partitions))
|
||||
else:
|
||||
pv = self.loopdev
|
||||
pvcreate(pv)
|
||||
vgcreate(vgname, pv)
|
||||
self.vgname = vgname
|
||||
for vol in volumes:
|
||||
lvcreate(vgname, vol['name'], vol.get('size'))
|
||||
self.volumes = volumes
|
||||
|
||||
def deactivate_vg(self):
|
||||
if not self.vgname:
|
||||
return
|
||||
run_cmd_logged('vgchange', '-an', self.vgname)
|
||||
|
||||
def make_filesystems(self):
|
||||
for mountpoint, dev, fstype, fsopts in self.filesystems:
|
||||
if not fstype:
|
||||
fstype = self.default_fstype
|
||||
log.info('Creating {} filesystem on {}'.format(fstype, dev))
|
||||
run_cmd('mkfs.{}'.format(fstype), dev)
|
||||
|
||||
def mount(self):
|
||||
self.mountpoint = tempfile.mkdtemp()
|
||||
for mountpoint, dev, fstype, fsopts in sorted(self.filesystems):
|
||||
path = os.path.join(self.mountpoint, mountpoint[1:])
|
||||
if not os.path.isdir(path):
|
||||
os.makedirs(path)
|
||||
log.info('Mounting {} on {}'.format(dev, path))
|
||||
run_cmd_logged('mount', dev, path)
|
||||
|
||||
def mount_extra(self):
|
||||
if self._extra_mounted:
|
||||
return
|
||||
for item in self.EXTRA_MOUNTS:
|
||||
mountpoint, dev, args = item[0], item[1], item[2:]
|
||||
path = os.path.join(self.mountpoint, mountpoint[1:])
|
||||
if not os.path.isdir(path):
|
||||
os.makedirs(path)
|
||||
run_cmd_logged('mount', dev, path, *args)
|
||||
|
||||
def unmount(self):
|
||||
if not self.mountpoint:
|
||||
return
|
||||
log.info('Unmounting {}'.format(self.mountpoint))
|
||||
run_cmd_logged('umount', '-R', self.mountpoint)
|
||||
os.rmdir(self.mountpoint)
|
||||
self.mountpoint = None
|
||||
self._mount_extra = False
|
||||
|
||||
def extract(self, filename):
|
||||
if not self.mountpoint:
|
||||
self.mount()
|
||||
log.info('Extracting {} to {}'.format(filename, self.mountpoint))
|
||||
run_cmd_logged('tar', '-xha', '--numeric-owner', '-f', filename,
|
||||
'-C', self.mountpoint)
|
||||
|
||||
def run_script(self, script, chroot=True):
|
||||
if not self._extra_mounted:
|
||||
self.mount_extra()
|
||||
name = os.path.basename(script)
|
||||
dest = os.path.join(self.mountpoint, 'tmp', name)
|
||||
log.debug('Copying {} to {}'.format(script, dest))
|
||||
shutil.copy(script, dest)
|
||||
os.chmod(dest, 0o0755)
|
||||
kwargs = {
|
||||
'env': self.SCRIPT_ENV.copy(),
|
||||
}
|
||||
if chroot:
|
||||
cmd = ('chroot', self.mountpoint, os.path.join('/tmp', name))
|
||||
else:
|
||||
cmd = (dest,)
|
||||
kwargs['env']['IMAGE_ROOT'] = self.mountpoint
|
||||
kwargs['cwd'] = self.mountpoint
|
||||
log.info('Running script: {}'.format(name))
|
||||
try:
|
||||
run_cmd_logged(*cmd, **kwargs)
|
||||
finally:
|
||||
os.unlink(dest)
|
||||
|
||||
def write_fstab(self, tmpfstmp=True):
|
||||
tmpl = '{dev}\t{mountpoint}\t{fstype}\t{fsopts}\t0 {passno}\n'
|
||||
log.info('Writing fstab')
|
||||
with open(os.path.join(self.mountpoint, 'etc/fstab'), 'w') as fstab:
|
||||
for mountpoint, dev, fstype, fsopts in sorted(self.filesystems):
|
||||
line = tmpl.format(
|
||||
dev=dev,
|
||||
mountpoint=mountpoint,
|
||||
fstype=fstype or self.default_fstype,
|
||||
fsopts=fsopts,
|
||||
passno=1 if mountpoint == '/' else 2,
|
||||
)
|
||||
log.debug(line.rstrip('\n'))
|
||||
fstab.write(line)
|
||||
if tmpfstmp:
|
||||
line = tmpl.format(
|
||||
dev='tmpfs',
|
||||
mountpoint='/tmp',
|
||||
fstype='tmpfs',
|
||||
fsopts='defaults',
|
||||
passno=0,
|
||||
)
|
||||
log.debug(line.rstrip('\n'))
|
||||
fstab.write(line)
|
||||
|
||||
|
||||
def gen_vgname(length=8):
|
||||
name = random.choice(string.ascii_lowercase)
|
||||
while len(name) < length - 1:
|
||||
name += random.choice(string.ascii_letters + string.digits)
|
||||
name += random.choice(string.ascii_lowercase)
|
||||
return name
|
||||
|
||||
|
||||
def inject_ssh_keys(root, ssh_keys):
|
||||
ssh_dir = os.path.join(root, 'root', '.ssh')
|
||||
if not os.path.isdir(ssh_dir):
|
||||
os.makedirs(ssh_dir)
|
||||
ssh_keys = set(ssh_keys)
|
||||
try:
|
||||
ssh_keys.remove(None)
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
for path in TRY_SSH_KEYS:
|
||||
if os.path.exists(path):
|
||||
ssh_keys.append(path)
|
||||
break
|
||||
with open(os.path.join(ssh_dir, 'authorized_keys'), 'a') as dest:
|
||||
for keyfile in ssh_keys:
|
||||
with open(keyfile) as src:
|
||||
pubkey = src.read().strip()
|
||||
log.info('Injecting SSH public key: {}'.format(keyfile))
|
||||
dest.write(pubkey + '\n')
|
||||
|
||||
|
||||
def lvcreate(vg_name, name, size=None):
|
||||
cmd = ['lvcreate', '-n', name]
|
||||
if size is not None:
|
||||
if '%' in size:
|
||||
cmd += ('-l', size)
|
||||
else:
|
||||
cmd += ('-L', size)
|
||||
else:
|
||||
cmd += ('-l', '100%FREE')
|
||||
cmd.append(vg_name)
|
||||
run_cmd_logged(*cmd)
|
||||
|
||||
|
||||
def parse_volumes(volstrings):
|
||||
pat = re.compile(r'(?<!\\)=')
|
||||
keys = ('size', 'fstype', 'name', 'fsopts', 'type')
|
||||
volumes = []
|
||||
for vol in volstrings:
|
||||
parts = pat.split(vol, 1)
|
||||
mountpoint = parts[0]
|
||||
try:
|
||||
pairs = itertools.zip_longest(keys, parts[1].split(';'))
|
||||
params = dict((k, v) for k, v in pairs if v not in ('', None))
|
||||
except IndexError:
|
||||
params = {}
|
||||
params['mountpoint'] = mountpoint
|
||||
if 'name' not in params:
|
||||
if mountpoint.startswith('/'):
|
||||
name = mountpoint[1:].replace('/', '-')
|
||||
else:
|
||||
name = mountpoint
|
||||
if name == '':
|
||||
name = 'root'
|
||||
params['name'] = name
|
||||
if mountpoint == '/boot':
|
||||
volumes.insert(0, (mountpoint, params))
|
||||
elif mountpoint == '/':
|
||||
volumes.insert(1, (mountpoint, params))
|
||||
else:
|
||||
volumes.append((mountpoint, params))
|
||||
return collections.OrderedDict(volumes)
|
||||
|
||||
|
||||
def parse_size(size):
|
||||
if isinstance(size, int):
|
||||
return size
|
||||
|
||||
m = SIZE_RE.match(size.lower())
|
||||
if not m:
|
||||
raise ValueError('Invalid size: {}'.format(size))
|
||||
parts = m.groupdict()
|
||||
if parts['unit'] in (None, 'b', 'byte', 'bytes'):
|
||||
factor = 1
|
||||
else:
|
||||
try:
|
||||
factor = UNITS[parts['unit']]
|
||||
except KeyError:
|
||||
raise ValueError('Invalid size : {}'.format(size))
|
||||
return int(float(parts['value']) * factor)
|
||||
|
||||
|
||||
def pvcreate(*pvs):
|
||||
run_cmd_logged('pvcreate', *pvs)
|
||||
|
||||
|
||||
def pvscan(dev=None):
|
||||
cmd = ['pvscan']
|
||||
if dev:
|
||||
cmd += ('--cache', dev)
|
||||
run_cmd_logged(*cmd)
|
||||
|
||||
|
||||
def run_cmd(*args, **kwargs):
|
||||
cmd = list(args)
|
||||
log.debug('EXEC {}'.format(' '.join(cmd)))
|
||||
with open(os.devnull) as nul:
|
||||
p = subprocess.Popen(cmd, stdin=nul, stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE, **kwargs)
|
||||
out, err = p.communicate()
|
||||
if p.returncode != 0:
|
||||
msg = err.strip().decode()
|
||||
if not msg:
|
||||
msg = 'Command exited with status {}'.format(p.returncode)
|
||||
log.error(msg)
|
||||
raise CommandError(msg)
|
||||
else:
|
||||
return out.strip().decode()
|
||||
|
||||
|
||||
def run_cmd_logged(*args, **kwargs):
|
||||
for line in run_cmd(*args, **kwargs).splitlines():
|
||||
log.info(line)
|
||||
|
||||
|
||||
def vgcreate(name, *pvs):
|
||||
run_cmd_logged('vgcreate', name, *pvs)
|
||||
|
||||
|
||||
|
||||
def setup_logging(verbose=0):
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.DEBUG)
|
||||
stderr_handler = logging.StreamHandler()
|
||||
stderr_handler.setLevel(logging.ERROR)
|
||||
logger.addHandler(stderr_handler)
|
||||
stdout_handler = logging.StreamHandler(sys.stdout)
|
||||
stdout_handler.addFilter(MaxLevelFilter())
|
||||
logger.addHandler(stdout_handler)
|
||||
if verbose < 1:
|
||||
stdout_handler.setLevel(logging.WARNING)
|
||||
elif verbose < 2:
|
||||
stdout_handler.setLevel(logging.INFO)
|
||||
else:
|
||||
stdout_handler.setLevel(logging.DEBUG)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--verbose', '-v',
|
||||
action='count',
|
||||
default=0,
|
||||
help='Print additional status information',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--size', '-s',
|
||||
type=parse_size,
|
||||
default='3G',
|
||||
help='Disk image size',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--lvm', '-L',
|
||||
metavar='VGNAME',
|
||||
nargs='?',
|
||||
help='Create volumes using LVM, in the volume group VGNAME. If a '
|
||||
'volume group name is not specified, a unique name will be '
|
||||
'generated.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--no-lvm',
|
||||
dest='lvm',
|
||||
action='store_false',
|
||||
help='Do not use LVM',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--volume', '-V',
|
||||
metavar='VOL',
|
||||
dest='volumes',
|
||||
action='append',
|
||||
help='Define a new volume, using the format '
|
||||
'<mountpoint>[=<size>[;<fstype>[;<name>[;<fsopts>'
|
||||
'[;<parttype>]]]]]',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--default-fstype',
|
||||
metavar='FSTYPE',
|
||||
help='Default filesystem type to use for volumes that do not specify '
|
||||
'one of their own',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--no-tmpfs-tmp',
|
||||
dest='tmpfstmp',
|
||||
action='store_false',
|
||||
default=True,
|
||||
help='Do not write an fstab entry for /tmp on tmpfs',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--overlay', '-O',
|
||||
metavar='FILENAME',
|
||||
dest='overlays',
|
||||
action='append',
|
||||
default=[],
|
||||
help='Extract the contents of FILENAME onto the image',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--script', '-S',
|
||||
metavar='FILENAME',
|
||||
dest='scripts',
|
||||
action='append',
|
||||
default=[],
|
||||
help='Run a script inside the image',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--no-chroot',
|
||||
dest='chroot',
|
||||
action='store_false',
|
||||
default=True,
|
||||
help='Do not chroot into the image to run scripts - DANGEROUS!',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--inject-ssh-key', '-i',
|
||||
metavar='FILENAME',
|
||||
nargs='?',
|
||||
dest='inject_ssh_keys',
|
||||
action='append',
|
||||
help='Pre-authorize SSH public keys for root',
|
||||
)
|
||||
parser.add_argument(
|
||||
'stagetbz',
|
||||
help='Path to stage tarball',
|
||||
)
|
||||
parser.add_argument(
|
||||
'image',
|
||||
nargs='?',
|
||||
help='Path to destination image file',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.image is None:
|
||||
basename = os.path.splitext(args.stagetbz)[0]
|
||||
if basename.endswith('.tar'):
|
||||
basename = basename[:-4]
|
||||
args.image = basename + '.img'
|
||||
if args.volumes is None:
|
||||
args.volumes = ['/']
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
setup_logging(args.verbose)
|
||||
|
||||
if args.lvm is None:
|
||||
args.lvm = gen_vgname()
|
||||
|
||||
volumes = parse_volumes(args.volumes)
|
||||
if args.lvm:
|
||||
try:
|
||||
partitions = [volumes.pop('/boot'), LVM_PV_PART]
|
||||
except KeyError:
|
||||
partitions = None
|
||||
else:
|
||||
partitions = volumes.values()
|
||||
|
||||
try:
|
||||
image = Image.create(args.image)
|
||||
except OSError as e:
|
||||
sys.stderr.write('Failed to create image: {}\n'.format(e))
|
||||
raise SystemExit(os.EX_CANTCREAT)
|
||||
if args.default_fstype:
|
||||
image.default_fstype = args.default_fstype
|
||||
with image:
|
||||
image.resize(args.size)
|
||||
if partitions:
|
||||
image.partition(partitions)
|
||||
if args.lvm:
|
||||
image.setup_lvm(args.lvm, volumes.values())
|
||||
image.make_filesystems()
|
||||
image.extract(args.stagetbz)
|
||||
for overlay in args.overlays:
|
||||
image.extract(overlay)
|
||||
image.write_fstab(tmpfstmp=args.tmpfstmp)
|
||||
for script in args.scripts:
|
||||
image.run_script(script, chroot=args.chroot)
|
||||
if args.inject_ssh_keys:
|
||||
inject_ssh_keys(image.mountpoint, args.inject_ssh_keys)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Loading…
Reference in New Issue