<commit_before>#!/bin/env python

# 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 google
import struct
import os
import sys
import json
import pb2dict
import array

import magic
from pb import *

#
# 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):
		"""
		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
			pb = self.payload()
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			pb.ParseFromString(f.read(size))
			entry = pb2dict.pb2dict(pb, pretty)

			# Read extra
			if self.extra_handler:
				entry['extra'] = self.extra_handler.load(f, pb)

			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
			pb = self.payload()
			pb2dict.dict2pb(entry, pb)
			pb_str = pb.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, pb)

	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):
		entries = []

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

			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):
		pb = pagemap_head()
		for item in entries:
			pb2dict.dict2pb(item, pb)
			pb_str = pb.SerializeToString()
			size = len(pb_str)
			f.write(struct.pack('i', size))
			f.write(pb_str)

			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


# 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 data.encode('base64')

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

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

	def dump(self, extra, f, pb):
		data = extra.decode('base64')
		f.write(data)

class ghost_file_extra_handler:
	def load(self, f, pb):
		data = f.read()
		return data.encode('base64')

	def dump(self, extra, f, pb):
		data = extra.decode('base64')
		f.write(data)

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

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

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

		return d

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

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

class ipc_sem_set_handler:
	def load(self, f, pb):
		entry = pb2dict.pb2dict(pb)
		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.fromfile(f, entry['nsems'])
		f.seek(rounded - size, 1)
		return s.tolist()

	def dump(self, extra, f, pb):
		entry = pb2dict.pb2dict(pb)
		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")
		s.tofile(f)
		f.write('\0' * (rounded - size))

class ipc_msg_queue_handler:
	def load(self, f, pb):
		entry = pb2dict.pb2dict(pb)
		messages = []
		for x in range (0, entry['qnum']):
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			msg = 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(data.encode('base64'))
		return messages
<commit_msg>crit: Fix ipc load/dump for io.BytesIO input/output<commit_after>#!/bin/env python

# 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 google
import struct
import os
import sys
import json
import pb2dict
import array

import magic
from pb import *

#
# 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):
		"""
		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
			pb = self.payload()
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			pb.ParseFromString(f.read(size))
			entry = pb2dict.pb2dict(pb, pretty)

			# Read extra
			if self.extra_handler:
				entry['extra'] = self.extra_handler.load(f, pb)

			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
			pb = self.payload()
			pb2dict.dict2pb(entry, pb)
			pb_str = pb.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, pb)

	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):
		entries = []

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

			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):
		pb = pagemap_head()
		for item in entries:
			pb2dict.dict2pb(item, pb)
			pb_str = pb.SerializeToString()
			size = len(pb_str)
			f.write(struct.pack('i', size))
			f.write(pb_str)

			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


# 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 data.encode('base64')

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

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

	def dump(self, extra, f, pb):
		data = extra.decode('base64')
		f.write(data)

class ghost_file_extra_handler:
	def load(self, f, pb):
		data = f.read()
		return data.encode('base64')

	def dump(self, extra, f, pb):
		data = extra.decode('base64')
		f.write(data)

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

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

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

		return d

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

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

class ipc_sem_set_handler:
	def load(self, f, pb):
		entry = pb2dict.pb2dict(pb)
		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, pb):
		entry = pb2dict.pb2dict(pb)
		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))

class ipc_msg_queue_handler:
	def load(self, f, pb):
		entry = pb2dict.pb2dict(pb)
		messages = []
		for x in range (0, entry['qnum']):
			buf = f.read(4)
			if buf == '':
				break
			size, = struct.unpack('i', buf)
			msg = 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(data.encode('base64'))
		return messages

{'fraction_non_alphanumeric': 0.13067823895849373, 'fraction_numerical': 0.012878840904318612, 'mean_word_length': 3.488846999685831, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '320713', 'n_tokens_mistral': 5757, 'n_tokens_neox': 5162, 'n_words': 1561}