[CRIU] [PATCH 09/10] crit: add ability to convert to/from json, v2

[Ruslan Kuprieiev]

v2, don't fail if json is not in "nice" format

Signed-off-by: Ruslan Kuprieiev <kupruser at gmail.com>
---
 crit                     |  6 +++-
 pycriu/images/images.py  | 57 ++++++++++++++++++++++++++++++--
 pycriu/images/pb2dict.py | 84 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 144 insertions(+), 3 deletions(-)
 create mode 100644 pycriu/images/pb2dict.py

diff --git a/crit b/crit
index 438ef3b..657610a 100755
--- a/crit
+++ b/crit
@@ -20,6 +20,10 @@ def handle_cmdline_opts():
 			    '--format',
 			    choices = ['raw', 'nice'],
 			help = 'well-formated output (by default: raw for files and nice for stdout)')
+	parser.add_argument('-t',
+			    '--type',
+			    choices = ['json'],
+			help = 'Type of output file (by default img is converted to txt and vice versa)')
 
 	opts = vars(parser.parse_args())
 
@@ -52,7 +56,7 @@ def convert_img(opts):
 	# img remembers data type it was parsed from, so
 	# SerializeToString will return data in type that
 	# is oposite to the original.
-	out_str = img.SerializeToString(nice=nice)
+	out_str = img.SerializeToString(nice=nice, typ=opts['type'])
 
 	# If no output file is set -- stdout is used.
 	if opts['out']:
diff --git a/pycriu/images/images.py b/pycriu/images/images.py
index afcaa87..ee7c9cf 100644
--- a/pycriu/images/images.py
+++ b/pycriu/images/images.py
@@ -4,6 +4,8 @@ import google
 import struct
 import os
 import sys
+import json
+import pb2dict
 
 from magic import *
 from pb import *
@@ -153,10 +155,33 @@ class criu_image:
 			self.msgs.append(pb)
 			pb = self.desc.body()
 
+	def ParseFromStringJSON(self, s):
+		img = json.loads(s)
+		magic = img["magic"]
+
+		descs = criu_image_descs()
+		try:
+			self.desc = descs.get(magic)
+		except Exception as e:
+			raise e
+
+		pb = self.desc.head()
+		index = 0
+		while pb != None and img.get(str(index)) != None:
+			entry = img[str(index)]
+			entry.pop("pbtype", None)
+			pb2dict.dict2pb(entry, pb)
+			self.msgs.append(pb)
+			pb = self.desc.body()
+			index += 1
+
 	def ParseFromString(self, s):
 		if s[0] == '#':
 			self.orig_type = 'txt'
 			self.ParseFromStringTXT(s)
+		elif s[0] == '{':
+			self.orig_type = 'json'
+			self.ParseFromStringJSON(s)
 		else:
 			self.orig_type = 'img'
 			self.ParseFromStringIMG(s)
@@ -201,8 +226,36 @@ class criu_image:
 		s.seek(0)
 		return s.read()
 
-	def SerializeToString(self, nice=False):
-		if self.orig_type == 'img':
+	def SerializeToStringJSON(self, nice=False):
+		if nice:
+			magic = self.desc.magic_name
+			indent = 4
+		else:
+			magic = self.desc.magic
+			indent = 0
+
+		img = {}
+		img["magic"] = magic
+
+		index = 0
+		for pb in self.msgs:
+			js = pb2dict.pb2dict(pb)
+
+			if nice:
+				if index == 0:
+					js["pbtype"] = self.desc.head_name
+				else:
+					js["pbtype"] = self.desc.body_name
+
+			img[str(index)] = js
+			index += 1
+
+		return json.dumps(img, indent=indent)
+
+	def SerializeToString(self, nice=False, typ=None):
+		if typ == 'json':
+			return self.SerializeToStringJSON(nice=nice)
+		elif self.orig_type == 'img':
 			return self.SerializeToStringTXT(nice=nice)
 		else:
 			return self.SerializeToStringIMG()
diff --git a/pycriu/images/pb2dict.py b/pycriu/images/pb2dict.py
new file mode 100644
index 0000000..b304824
--- /dev/null
+++ b/pycriu/images/pb2dict.py
@@ -0,0 +1,84 @@
+import google
+import io
+
+# pb2dict and dict2pb are using protobuf text format to properly
+# convert protobuf msgs to/from dictionary.
+
+def pb2dict(pb):
+	"""
+	Convert protobuf msg to dictionary
+	"""
+	pb_text = io.BytesIO('')
+	google.protobuf.text_format.PrintMessage(pb, pb_text)
+	pb_text.seek(0)
+	return _text2dict(pb_text)
+
+def _text2dict(pb_text):
+	"""
+	Convert protobuf text format msg to dict
+	"""
+	d = {}
+	while True:
+		s = pb_text.readline()
+		s.strip()
+		if s == '' or '}' in s:
+			break
+
+		name, value = s.split()
+		if value == '{':
+			value = _text2dict(pb_text)
+		elif name.endswith(':'):
+			name = name[:-1]
+		else:
+			raise Exception("Unknown format" + s)
+
+		if d.get(name):
+			if not isinstance(d[name], list):
+				d[name] = [d[name]]
+			d[name].append(value)
+		else:
+			d[name] = value
+
+	return d
+
+def dict2pb(d, pb):
+	"""
+	Convert dictionary to protobuf msg
+	"""
+	pb_text = io.BytesIO('')
+	_dict2text(d, pb_text, 0)
+	pb_text.seek(0)
+	s = pb_text.read()
+	google.protobuf.text_format.Merge(s, pb)
+
+def _write_struct(name, text, indent, inside):
+	"""
+	Convert "inside" dict to protobuf text format
+	wrap it inside block named "name" and write
+	it to "text".
+	"""
+	text.write(indent*" " + name.encode() + " {\n")
+	_dict2text(inside, text, indent+2)
+	text.write(indent*" " + "}\n")
+
+def _write_field(name, value, text, indent):
+	"""
+	Write "name: value" to "text".
+	"""
+	text.write(indent*" " + name.encode() + ": " + value.encode() + "\n")
+
+def _dict2text(d, pb_text, indent):
+	"""
+	Convert dict to protobuf text format
+	"""
+	for name, value in d.iteritems():
+		if isinstance(value, unicode):
+			_write_field(name, value, pb_text, indent)
+		elif isinstance(value, list):
+			for x in value:
+				if isinstance(x, dict):
+					_write_struct(name, pb_text, indent, x)
+				else:
+					_write_field(name, x, pb_text, indent)
+		else:
+			_write_struct(name, pb_text, indent, value)
-- 
1.9.3