criu/lib/py/images/images.py

# This file contains methods to deal with criu images.
#
# According to http://criu.org/Images, criu images can be described
# with such IOW:
#
# IMAGE_FILE ::= MAGIC { ENTRY }
# ENTRY      ::= SIZE PAYLOAD [ EXTRA ]
# PAYLOAD    ::= "message encoded in ProtocolBuffer format"
# EXTRA      ::= "arbitrary blob, depends on the PAYLOAD contents"
#
# MAGIC      ::= "32 bit integer"
# SIZE       ::= "32 bit integer, equals the PAYLOAD length"
#
# Images v1.1 NOTE: MAGIC now consist of 2 32 bit integers, first one is
#	MAGIC_COMMON or MAGIC_SERVICE and the second one is same as MAGIC
#	in images V1.0. We don't keep "first" magic in json images.
#
# In order to convert images to human-readable format, we use dict(json).
# Using json not only allows us to easily read\write images, but also
# to use a great variety of tools out there to manipulate them.
# It also allows us to clearly describe criu images structure.
#
# Using dict(json) format, criu images can be described like:
#
# {
#	'magic' : 'FOO',
#	'entries' : [
#		entry,
#		...
#	]
# }
#
# Entry, in its turn, could be described as:
#
# {
#	pb_msg,
#	'extra' : extra_msg
# }
#
import io
import base64
import struct
import os
import array

from . import magic
from . import pb
from . import pb2dict

if "encodebytes" not in dir(base64):
	base64.encodebytes = base64.encodestring
	base64.decodebytes = base64.decodestring

#
# Predefined hardcoded constants
sizeof_u16 = 2
sizeof_u32 = 4
sizeof_u64 = 8

# A helper for rounding
def round_up(x,y):
	return (((x - 1) | (y - 1)) + 1)

class MagicException(Exception):
	def __init__(self, magic):
		self.magic = magic

# Generic class to handle loading/dumping criu images entries from/to bin
# format to/from dict(json).
class entry_handler:
	"""
	Generic class to handle loading/dumping criu images
	entries from/to bin format to/from dict(json).
	"""
	def __init__(self, payload, extra_handler=None):
		"""
		Sets payload class and extra handler class.
		"""
		self.payload		= payload
		self.extra_handler	= extra_handler

	def load(self, f, pretty = False, no_payload = False):
		"""
		Convert criu image entries from binary format to dict(json).
		Takes a file-like object and returnes a list with entries in
		dict(json) format.
		"""
		entries = []

		while True:
			entry = {}

			# Read payload
			pbuff = self.payload()
			buf = f.read(4)
			if buf == b'':
				break
			size, = struct.unpack('i', buf)
			pbuff.ParseFromString(f.read(size))
			entry = pb2dict.pb2dict(pbuff, pretty)

			# Read extra
			if self.extra_handler:
				if no_payload:
					def human_readable(num):
						for unit in ['','K','M','G','T','P','E','Z']:
							if num < 1024.0:
								if int(num) == num:
									return "%d%sB" % (num, unit)
								else:
									return "%.1f%sB" % (num, unit)
							num /= 1024.0
						return "%.1fYB" % num

					pl_size = self.extra_handler.skip(f, pbuff)
					entry['extra'] = '... <%s>' % human_readable(pl_size)
				else:
					entry['extra'] = self.extra_handler.load(f, pbuff)

			entries.append(entry)

		return entries

	def loads(self, s, pretty = False):
		"""
		Same as load(), but takes a string as an argument.
		"""
		f = io.BytesIO(s)
		return self.load(f, pretty)

	def dump(self, entries, f):
		"""
		Convert criu image entries from dict(json) format to binary.
		Takes a list of entries and a file-like object to write entries
		in binary format to.
		"""
		for entry in entries:
			extra = entry.pop('extra', None)

			# Write payload
			pbuff = self.payload()
			pb2dict.dict2pb(entry, pbuff)
			pb_str = pbuff.SerializeToString()
			size = len(pb_str)
			f.write(struct.pack('i', size))
			f.write(pb_str)

			# Write extra
			if self.extra_handler and extra:
				self.extra_handler.dump(extra, f, pbuff)

	def dumps(self, entries):
		"""
		Same as dump(), but doesn't take file-like object and just
		returns a string.
		"""
		f = io.BytesIO('')
		self.dump(entries, f)
		return f.read()

	def count(self, f):
		"""
		Counts the number of top-level object in the image file
		"""
		entries = 0

		while True:
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			f.seek(size, 1)
			entries += 1

		return entries

# Special handler for pagemap.img
class pagemap_handler:
	"""
	Special entry handler for pagemap.img, which is unique in a way
	that it has a header of pagemap_head type followed by entries
	of pagemap_entry type.
	"""
	def load(self, f, pretty = False, no_payload = False):
		entries = []

		pbuff = pb.pagemap_head()
		while True:
			buf = f.read(4)
			if buf == b'':
				break
			size, = struct.unpack('i', buf)
			pbuff.ParseFromString(f.read(size))
			entries.append(pb2dict.pb2dict(pbuff, pretty))

			pbuff = pb.pagemap_entry()

		return entries

	def loads(self, s, pretty = False):
		f = io.BytesIO(s)
		return self.load(f, pretty)

	def dump(self, entries, f):
		pbuff = pb.pagemap_head()
		for item in entries:
			pb2dict.dict2pb(item, pbuff)
			pb_str = pbuff.SerializeToString()
			size = len(pb_str)
			f.write(struct.pack('i', size))
			f.write(pb_str)

			pbuff = pb.pagemap_entry()

	def dumps(self, entries):
		f = io.BytesIO('')
		self.dump(entries, f)
		return f.read()

	def count(self, f):
		return entry_handler(None).count(f) - 1

# Special handler for ghost-file.img
class ghost_file_handler:
	def load(self, f, pretty = False, no_payload = False):
		entries = []

		gf = pb.ghost_file_entry()
		buf = f.read(4)
		size, = struct.unpack('i', buf)
		gf.ParseFromString(f.read(size))
		g_entry = pb2dict.pb2dict(gf, pretty)

		if gf.chunks:
			entries.append(g_entry)
			while True:
				gc = pb.ghost_chunk_entry()
				buf = f.read(4)
				if buf == '':
					break
				size, = struct.unpack('i', buf)
				gc.ParseFromString(f.read(size))
				entry = pb2dict.pb2dict(gc, pretty)
				if no_payload:
					f.seek(gc.len, os.SEEK_CUR)
				else:
					entry['extra'] = base64.encodebytes(f.read(gc.len))
				entries.append(entry)
		else:
			if no_payload:
				f.seek(0, os.SEEK_END)
			else:
				g_entry['extra'] = base64.encodebytes(f.read())
			entries.append(g_entry)

		return entries

	def loads(self, s, pretty = False):
		f = io.BytesIO(s)
		return self.load(f, pretty)

	def dump(self, entries, f):
		pbuff = pb.ghost_file_entry()
		item = entries.pop(0)
		pb2dict.dict2pb(item, pbuff)
		pb_str = pbuff.SerializeToString()
		size = len(pb_str)
		f.write(struct.pack('i', size))
		f.write(pb_str)

		if pbuff.chunks:
			for item in entries:
				pbuff = pb.ghost_chunk_entry()
				pb2dict.dict2pb(item, pbuff)
				pb_str = pbuff.SerializeToString()
				size = len(pb_str)
				f.write(struct.pack('i', size))
				f.write(pb_str)
				f.write(base64.decodebytes(item['extra']))
		else:
			f.write(base64.decodebytes(item['extra']))

	def dumps(self, entries):
		f = io.BytesIO('')
		self.dump(entries, f)
		return f.read()


# In following extra handlers we use base64 encoding
# to store binary data. Even though, the nature
# of base64 is that it increases the total size,
# it doesn't really matter, because our images
# do not store big amounts of binary data. They
# are negligible comparing to pages size.
class pipes_data_extra_handler:
	def load(self, f, pload):
		size = pload.bytes
		data = f.read(size)
		return base64.encodebytes(data)

	def dump(self, extra, f, pload):
		data = base64.decodebytes(extra)
		f.write(data)

	def skip(self, f, pload):
		f.seek(pload.bytes, os.SEEK_CUR)
		return pload.bytes

class sk_queues_extra_handler:
	def load(self, f, pload):
		size = pload.length
		data = f.read(size)
		return base64.encodebytes(data)

	def dump(self, extra, f, _unused):
		data = base64.decodebytes(extra)
		f.write(data)

	def skip(self, f, pload):
		f.seek(pload.length, os.SEEK_CUR)
		return pload.length


class tcp_stream_extra_handler:
	def load(self, f, pbuff):
		d = {}

		inq	= f.read(pbuff.inq_len)
		outq	= f.read(pbuff.outq_len)

		d['inq']	= base64.encodebytes(inq)
		d['outq']	= base64.encodebytes(outq)

		return d

	def dump(self, extra, f, _unused):
		inq	= base64.decodebytes(extra['inq'])
		outq	= base64.decodebytes(extra['outq'])

		f.write(inq)
		f.write(outq)

	def skip(self, f, pbuff):
		f.seek(0, os.SEEK_END)
		return pbuff.inq_len + pbuff.outq_len

class ipc_sem_set_handler:
	def load(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = sizeof_u16 * entry['nsems']
		rounded = round_up(size, sizeof_u64)
		s = array.array('H')
		if s.itemsize != sizeof_u16:
			raise Exception("Array size mismatch")
		s.fromstring(f.read(size))
		f.seek(rounded - size, 1)
		return s.tolist()

	def dump(self, extra, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = sizeof_u16 * entry['nsems']
		rounded = round_up(size, sizeof_u64)
		s = array.array('H')
		if s.itemsize != sizeof_u16:
			raise Exception("Array size mismatch")
		s.fromlist(extra)
		if len(s) != entry['nsems']:
			raise Exception("Number of semaphores mismatch")
		f.write(s.tostring())
		f.write('\0' * (rounded - size))

	def skip(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = sizeof_u16 * entry['nsems']
		f.seek(round_up(size, sizeof_u64), os.SEEK_CUR)
		return size

class ipc_msg_queue_handler:
	def load(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		messages = []
		for x in range (0, entry['qnum']):
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			msg = pb.ipc_msg()
			msg.ParseFromString(f.read(size))
			rounded = round_up(msg.msize, sizeof_u64)
			data = f.read(msg.msize)
			f.seek(rounded - msg.msize, 1)
			messages.append(pb2dict.pb2dict(msg))
			messages.append(base64.encodebytes(data))
		return messages

	def dump(self, extra, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		for i in range (0, len(extra), 2):
			msg = pb.ipc_msg()
			pb2dict.dict2pb(extra[i], msg)
			msg_str = msg.SerializeToString()
			size = len(msg_str)
			f.write(struct.pack('i', size))
			f.write(msg_str)
			rounded = round_up(msg.msize, sizeof_u64)
			data = base64.decodebytes(extra[i + 1])
			f.write(data[:msg.msize])
			f.write('\0' * (rounded - msg.msize))

	def skip(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		pl_len = 0
		for x in range (0, entry['qnum']):
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			msg = pb.ipc_msg()
			msg.ParseFromString(f.read(size))
			rounded = round_up(msg.msize, sizeof_u64)
			f.seek(rounded, os.SEEK_CUR)
			pl_len += size + msg.msize

		return pl_len

class ipc_shm_handler:
	def load(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = entry['size']
		data = f.read(size)
		rounded = round_up(size, sizeof_u32)
		f.seek(rounded - size, 1)
		return base64.encodebytes(data)

	def dump(self, extra, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = entry['size']
		data = base64.decodebytes(extra)
		rounded = round_up(size, sizeof_u32)
		f.write(data[:size])
		f.write('\0' * (rounded - size))

	def skip(self, f, pbuff):
		entry = pb2dict.pb2dict(pbuff)
		size = entry['size']
		rounded = round_up(size, sizeof_u32)
		f.seek(rounded, os.SEEK_CUR)
		return size


handlers = {
	'INVENTORY'		: entry_handler(pb.inventory_entry),
	'CORE'			: entry_handler(pb.core_entry),
	'IDS'			: entry_handler(pb.task_kobj_ids_entry),
	'CREDS'			: entry_handler(pb.creds_entry),
	'UTSNS'			: entry_handler(pb.utsns_entry),
	'IPC_VAR'		: entry_handler(pb.ipc_var_entry),
	'FS'			: entry_handler(pb.fs_entry),
	'GHOST_FILE'		: ghost_file_handler(),
	'MM'			: entry_handler(pb.mm_entry),
	'CGROUP'		: entry_handler(pb.cgroup_entry),
	'TCP_STREAM'		: entry_handler(pb.tcp_stream_entry, tcp_stream_extra_handler()),
	'STATS'			: entry_handler(pb.stats_entry),
	'PAGEMAP'		: pagemap_handler(), # Special one
	'PSTREE'		: entry_handler(pb.pstree_entry),
	'REG_FILES'		: entry_handler(pb.reg_file_entry),
	'NS_FILES'		: entry_handler(pb.ns_file_entry),
	'EVENTFD_FILE'		: entry_handler(pb.eventfd_file_entry),
	'EVENTPOLL_FILE'	: entry_handler(pb.eventpoll_file_entry),
	'EVENTPOLL_TFD'		: entry_handler(pb.eventpoll_tfd_entry),
	'SIGNALFD'		: entry_handler(pb.signalfd_entry),
	'TIMERFD'		: entry_handler(pb.timerfd_entry),
	'INOTIFY_FILE'		: entry_handler(pb.inotify_file_entry),
	'INOTIFY_WD'		: entry_handler(pb.inotify_wd_entry),
	'FANOTIFY_FILE'		: entry_handler(pb.fanotify_file_entry),
	'FANOTIFY_MARK'		: entry_handler(pb.fanotify_mark_entry),
	'VMAS'			: entry_handler(pb.vma_entry),
	'PIPES'			: entry_handler(pb.pipe_entry),
	'FIFO'			: entry_handler(pb.fifo_entry),
	'SIGACT'		: entry_handler(pb.sa_entry),
	'NETLINK_SK'		: entry_handler(pb.netlink_sk_entry),
	'REMAP_FPATH'		: entry_handler(pb.remap_file_path_entry),
	'MNTS'			: entry_handler(pb.mnt_entry),
	'TTY_FILES'		: entry_handler(pb.tty_file_entry),
	'TTY_INFO'		: entry_handler(pb.tty_info_entry),
	'TTY_DATA'		: entry_handler(pb.tty_data_entry),
	'RLIMIT'		: entry_handler(pb.rlimit_entry),
	'TUNFILE'		: entry_handler(pb.tunfile_entry),
	'EXT_FILES'		: entry_handler(pb.ext_file_entry),
	'IRMAP_CACHE'		: entry_handler(pb.irmap_cache_entry),
	'FILE_LOCKS'		: entry_handler(pb.file_lock_entry),
	'FDINFO'		: entry_handler(pb.fdinfo_entry),
	'UNIXSK'		: entry_handler(pb.unix_sk_entry),
	'INETSK'		: entry_handler(pb.inet_sk_entry),
	'PACKETSK'		: entry_handler(pb.packet_sock_entry),
	'ITIMERS'		: entry_handler(pb.itimer_entry),
	'POSIX_TIMERS'		: entry_handler(pb.posix_timer_entry),
	'NETDEV'		: entry_handler(pb.net_device_entry),
	'PIPES_DATA'		: entry_handler(pb.pipe_data_entry, pipes_data_extra_handler()),
	'FIFO_DATA'		: entry_handler(pb.pipe_data_entry, pipes_data_extra_handler()),
	'SK_QUEUES'		: entry_handler(pb.sk_packet_entry, sk_queues_extra_handler()),
	'IPCNS_SHM'		: entry_handler(pb.ipc_shm_entry, ipc_shm_handler()),
	'IPCNS_SEM'		: entry_handler(pb.ipc_sem_entry, ipc_sem_set_handler()),
	'IPCNS_MSG'		: entry_handler(pb.ipc_msg_entry, ipc_msg_queue_handler()),
	'NETNS'			: entry_handler(pb.netns_entry),
	'USERNS'		: entry_handler(pb.userns_entry),
	'SECCOMP'		: entry_handler(pb.seccomp_entry),
	'AUTOFS'		: entry_handler(pb.autofs_entry),
	'FILES'                 : entry_handler(pb.file_entry),
	'CPUINFO'		: entry_handler(pb.cpuinfo_entry),
	}

def __rhandler(f):
	# Images v1.1 NOTE: First read "first" magic.
	img_magic, = struct.unpack('i', f.read(4))
	if img_magic in (magic.by_name['IMG_COMMON'], magic.by_name['IMG_SERVICE']):
		img_magic, = struct.unpack('i', f.read(4))

	try:
		m = magic.by_val[img_magic]
	except:
		raise MagicException(img_magic)

	try:
		handler = handlers[m]
	except:
		raise Exception("No handler found for image with magic " + m)

	return m, handler

def load(f, pretty = False, no_payload = False):
	"""
	Convert criu image from binary format to dict(json).
	Takes a file-like object to read criu image from.
	Returns criu image in dict(json) format.
	"""
	image = {}

	m, handler = __rhandler(f)

	image['magic'] = m
	image['entries'] = handler.load(f, pretty, no_payload)

	return image

def info(f):
	res = {}

	m, handler = __rhandler(f)

	res['magic'] = m
	res['count'] = handler.count(f)

	return res

def loads(s, pretty = False):
	"""
	Same as load(), but takes a string.
	"""
	f = io.BytesIO(s)
	return load(f, pretty)

def dump(img, f):
	"""
	Convert criu image from dict(json) format to binary.
	Takes an image in dict(json) format and file-like
	object to write to.
	"""
	m = img['magic']
	magic_val = magic.by_name[img['magic']]

	# Images v1.1 NOTE: use "second" magic to identify what "first"
	# should be written.
	if m != 'INVENTORY':
		if m in ('STATS', 'IRMAP_CACHE'):
			f.write(struct.pack('i', magic.by_name['IMG_SERVICE']))
		else:
			f.write(struct.pack('i', magic.by_name['IMG_COMMON']))

	f.write(struct.pack('i', magic_val))

	try:
		handler = handlers[m]
	except:
		raise Exception("No handler found for image with such magic")

	handler.dump(img['entries'], f)

def dumps(img):
	"""
	Same as dump(), but takes only an image and returns
	a string.
	"""
	f = io.BytesIO(b'')
	dump(img, f)
	return f.getvalue()