Commit d65c4ad6 authored by Constantinos Venetsanopoulos's avatar Constantinos Venetsanopoulos
Browse files

Merge PR #16 (feature-code-cleanup) from skalkoto/snf-image-creator

Small fixes and code beautification.
parents 78aac79a 12f4aad8
......@@ -61,8 +61,11 @@ def create_image(d, media, out, tmp):
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
snapshot = disk.snapshot()
image = disk.get_image(snapshot)
# There is no need to snapshot the media if it was created by the Disk
# instance as a temporary object.
device = disk.device if disk.source == '/' else disk.snapshot()
image = disk.get_image(device)
out.output("Collecting image metadata ...")
metadata = {}
......
......@@ -353,11 +353,11 @@ def create_image(session):
try:
out.clear()
#Sysprep
# Sysprep
image.os.do_sysprep()
metadata = image.os.meta
#Shrink
# Shrink
size = image.shrink()
session['shrinked'] = True
update_background_title(session)
......@@ -365,7 +365,7 @@ def create_image(session):
metadata.update(image.meta)
metadata['DESCRIPTION'] = wizard['ImageDescription']
#MD5
# MD5
md5 = MD5(out)
session['checksum'] = md5.compute(image.device, size)
......
......@@ -72,6 +72,7 @@ class Disk(object):
"""
self._cleanup_jobs = []
self._images = []
self._device = None
self.source = source
self.out = output
self.meta = {}
......@@ -123,29 +124,36 @@ class Disk(object):
job, args = self._cleanup_jobs.pop()
job(*args)
def snapshot(self):
"""Creates a snapshot of the original source media of the Disk
instance.
"""
@property
def device(self):
"""Convert the source media into a block device"""
if self._device is not None:
return self._device
self.out.output("Examining source media `%s' ..." % self.source, False)
sourcedev = self.source
mode = os.stat(self.source).st_mode
if stat.S_ISDIR(mode):
self.out.success('looks like a directory')
return self._dir_to_disk()
self._device = self._dir_to_disk()
elif stat.S_ISREG(mode):
self.out.success('looks like an image file')
sourcedev = self._losetup(self.source)
self._device = self._losetup(self.source)
elif not stat.S_ISBLK(mode):
raise FatalError("Invalid media source. Only block devices, "
"regular files and directories are supported.")
else:
self.out.success('looks like a block device')
self._device = self.source
# Take a snapshot and return it to the user
return self._device
def snapshot(self):
"""Creates a snapshot of the original source media of the Disk
instance.
"""
self.out.output("Snapshotting media source ...", False)
size = blockdev('--getsz', sourcedev)
size = blockdev('--getsz', self.device)
cowfd, cow = tempfile.mkstemp(dir=self.tmp)
os.close(cowfd)
self._add_cleanup(os.unlink, cow)
......@@ -158,7 +166,7 @@ class Disk(object):
try:
try:
os.write(tablefd, "0 %d snapshot %s %s n 8" %
(int(size), sourcedev, cowdev))
(int(size), self.device, cowdev))
finally:
os.close(tablefd)
......
......@@ -267,7 +267,7 @@ class GPTPartitionTable(object):
lba_count = new_size // BLOCKSIZE
# Correct MBR
#TODO: Check if the partition tables is hybrid
# TODO: Check if the partition tables is hybrid
self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
# Fix Primary header
......
......@@ -137,8 +137,8 @@ class Image(object):
else:
self.g.set_recovery_proc(0)
#self.g.set_trace(1)
#self.g.set_verbose(1)
# self.g.set_trace(1)
# self.g.set_verbose(1)
self.out.output('Launching helper VM (may take a while) ...', False)
# self.progressbar = self.out.Progress(100, "Launching helper VM",
......@@ -180,7 +180,8 @@ class Image(object):
self.guestfs_enabled = False
self.out.success('done')
def _get_os(self):
@property
def os(self):
"""Return an OS class instance for this image"""
if hasattr(self, "_os"):
return self._os
......@@ -195,8 +196,6 @@ class Image(object):
return self._os
os = property(_get_os)
def destroy(self):
"""Destroy this Image instance."""
......
......@@ -257,9 +257,10 @@ def image_creator():
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
snapshot = disk.snapshot()
image = disk.get_image(snapshot, sysprep_params=options.sysprep_params)
# There is no need to snapshot the media if it was created by the Disk
# instance as a temporary object.
device = disk.device if disk.source == '/' else disk.snapshot()
image = disk.get_image(device, sysprep_params=options.sysprep_params)
if image.is_unsupported() and not options.allow_unsupported:
raise FatalError(
......@@ -323,7 +324,7 @@ def image_creator():
os.path.basename(options.outfile)))
out.success('done')
# Destroy the image instance. We only need the snapshot from now on
# Destroy the image instance. We only need the disk device from now on
disk.destroy_image(image)
out.output()
......@@ -331,7 +332,7 @@ def image_creator():
uploaded_obj = ""
if options.upload:
out.output("Uploading image to the storage service:")
with open(snapshot, 'rb') as f:
with open(device, 'rb') as f:
uploaded_obj = kamaki.upload(
f, size, options.upload,
"(1/3) Calculating block hashes",
......
......@@ -82,29 +82,25 @@ def sysprep(message, enabled=True, **kwargs):
class SysprepParam(object):
"""This class represents an system preparation parameter"""
"""This class represents a system preparation parameter"""
def __init__(self, type, default, description):
def __init__(self, type, default, description, check=lambda x: x):
type_checker = {"posint": self._check_posint,
"string": self._check_string,
"file": self._check_fname,
"dir": self._check_dname}
assert type in type_checker.keys(), "Invalid parameter type: %s" % type
assert hasattr(self, "_check_%s" % type), "Invalid type: %s" % type
self.type = type
self.default = default
self.description = description
self.value = default
self.error = None
self._checker = type_checker[type]
self.check = check
def set_value(self, value):
"""Update the value of the parameter"""
check_type = getattr(self, "_check_%s" % self.type)
try:
self.value = self._checker(value)
self.value = self.check(check_type(value))
except ValueError as e:
self.error = e.message
return False
......@@ -126,7 +122,7 @@ class SysprepParam(object):
"""Check if a value is a string"""
return str(value)
def _check_fname(self, value):
def _check_file(self, value):
"""Check if the value is a valid filename"""
value = str(value)
......@@ -146,7 +142,7 @@ class SysprepParam(object):
raise ValueError("Invalid filename")
def _check_dname(self, value):
def _check_dir(self, value):
"""Check if the value is a valid directory"""
value = str(value)
......@@ -160,7 +156,7 @@ class SysprepParam(object):
raise ValueError("Invalid dirname")
def add_sysprep_param(name, type, default, descr):
def add_sysprep_param(name, type, default, descr, check=lambda x: x):
"""Decorator for __init__ that adds the definition for a system preparation
parameter in an instance of an os_type class
"""
......@@ -171,8 +167,8 @@ def add_sysprep_param(name, type, default, descr):
if not hasattr(self, 'sysprep_params'):
self.sysprep_params = {}
self.sysprep_params[name] = SysprepParam(type, default, descr)
self.sysprep_params[name] = SysprepParam(type, default, descr,
check)
init(self, *args, **kwargs)
return inner
return wrapper
......
......@@ -55,7 +55,7 @@ class Freebsd(Unix):
super(Freebsd, self)._do_collect_metadata()
self.meta["USERS"] = " ".join(self._get_passworded_users())
#The original product name key is long and ugly
# The original product name key is long and ugly
self.meta['DESCRIPTION'] = \
self.meta['DESCRIPTION'].split('#')[0].strip()
......
......@@ -33,6 +33,10 @@ WINDOWS_SETUP_STATES = (
"IMAGE_STATE_SPECIALIZE_RESEAL_TO_OOBE",
"IMAGE_STATE_SPECIALIZE_RESEAL_TO_AUDIT")
REG_SZ = lambda k, v: {'key': k, 't': 1L,
'value': (v + '\x00').encode('utf-16le')}
REG_BINARY = lambda k, v: {'key': k, 't': 3L, 'value': v}
REG_DWORD = lambda k, v: {'key': k, 't': 4L, 'value': struct.pack('<I', v)}
class Registry(object):
"""Windows Registry manipulation methods"""
......@@ -148,12 +152,11 @@ class Registry(object):
for desc, cmd in commands.items():
assert type(desc) is str and type(cmd) is str
value = {'key': desc, 't': 1, 'value': cmd.encode('utf-16le')}
hive.node_set_value(runonce, value)
hive.node_set_value(runonce, REG_SZ(desc, cmd))
hive.commit(None)
def enable_autologon(self, username, password="", autoadminlogon=True):
def enable_autologon(self, username, password=""):
"""Enable automatic logon for a specific user"""
assert type(username) is str and type(password) is str
......@@ -165,16 +168,9 @@ class Registry(object):
'Winlogon'):
winlogon = hive.node_get_child(winlogon, child)
hive.node_set_value(winlogon,
{'key': 'DefaultUserName', 't': 1,
'value': username.encode('utf-16le')})
hive.node_set_value(winlogon,
{'key': 'DefaultPassword', 't': 1,
'value': password.encode('utf-16le')})
hive.node_set_value(
winlogon,
{'key': 'AutoAdminLogon', 't': 1,
'value': ("%d" % int(autoadminlogon)).encode('utf-16le')})
hive.node_set_value(winlogon, REG_SZ('DefaultUserName', username))
hive.node_set_value(winlogon, REG_SZ('DefaultPassword', password))
hive.node_set_value(winlogon, REG_SZ('AutoAdminLogon', "1"))
hive.commit(None)
......@@ -212,9 +208,8 @@ class Registry(object):
assert hive.value_type(old_value)[1] == 4
old_values.append(hive.value_dword(old_value))
hive.node_set_value(
node, {'key': 'EnableFirewall', 't': 4L,
'value': struct.pack("<I", new_values.pop(0))})
hive.node_set_value(node, REG_DWORD('EnableFirewall',
new_values.pop(0)))
hive.commit(None)
return old_values
......@@ -255,10 +250,8 @@ class Registry(object):
elif value == 0:
return False
new_value = {'key': "LocalAccountTokenFilterPolicy", 't': 4L,
'value': struct.pack("<I", value)}
hive.node_set_value(key, new_value)
hive.node_set_value(
key, REG_DWORD("LocalAccountTokenFilterPolicy", value))
hive.commit(None)
return True
......@@ -326,7 +319,7 @@ class Registry(object):
fmt = '%ds4x8s4x%ds' % (0xa0, len(v_val) - 0xb0)
new = ("\x00" * 4).join(struct.unpack(fmt, v_val))
hive.node_set_value(rid_node, {'key': "V", 't': 3L, 'value': new})
hive.node_set_value(rid_node, REG_BINARY('V', new))
hive.commit(None)
parent['old'] = v_val
......
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