| Current File : //usr/lib/python2.7/site-packages/pycriu/images/pb2dict.py |
from google.protobuf.descriptor import FieldDescriptor as FD
import opts_pb2
from ipaddress import IPv4Address, ip_address
from ipaddress import IPv6Address
import socket
import collections
import os
import base64
import quopri
if "encodebytes" not in dir(base64):
base64.encodebytes = base64.encodestring
base64.decodebytes = base64.decodestring
# pb2dict and dict2pb are methods to convert pb to/from dict.
# Inspired by:
# protobuf-to-dict - https://github.com/benhodgson/protobuf-to-dict
# protobuf-json - https://code.google.com/p/protobuf-json/
# protobuf source - https://code.google.com/p/protobuf/
# Both protobuf-to-dict/json do not fit here because of several reasons,
# here are some of them:
# - both have a common bug in treating optional field with empty
# repeated inside.
# - protobuf-to-json is not avalible in pip or in any other python
# repo, so it is hard to distribute and we can't rely on it.
# - both do not treat enums in a way we would like to. They convert
# protobuf enum to int, but we need a string here, because it is
# much more informative. BTW, protobuf text_format converts pb
# enums to string value too. (i.e. "march : x86_64" is better then
# "march : 1").
_basic_cast = {
FD.TYPE_FIXED64: int,
FD.TYPE_FIXED32: int,
FD.TYPE_SFIXED64: int,
FD.TYPE_SFIXED32: int,
FD.TYPE_INT64: int,
FD.TYPE_UINT64: int,
FD.TYPE_SINT64: int,
FD.TYPE_INT32: int,
FD.TYPE_UINT32: int,
FD.TYPE_SINT32: int,
FD.TYPE_BOOL: bool,
FD.TYPE_STRING: str
}
def _marked_as_hex(field):
return field.GetOptions().Extensions[opts_pb2.criu].hex
def _marked_as_ip(field):
return field.GetOptions().Extensions[opts_pb2.criu].ipadd
def _marked_as_flags(field):
return field.GetOptions().Extensions[opts_pb2.criu].flags
def _marked_as_dev(field):
return field.GetOptions().Extensions[opts_pb2.criu].dev
def _marked_as_odev(field):
return field.GetOptions().Extensions[opts_pb2.criu].odev
def _marked_as_dict(field):
return field.GetOptions().Extensions[opts_pb2.criu].dict
def _custom_conv(field):
return field.GetOptions().Extensions[opts_pb2.criu].conv
mmap_prot_map = [
('PROT_READ', 0x1),
('PROT_WRITE', 0x2),
('PROT_EXEC', 0x4),
]
mmap_flags_map = [
('MAP_SHARED', 0x1),
('MAP_PRIVATE', 0x2),
('MAP_ANON', 0x20),
('MAP_GROWSDOWN', 0x0100),
]
mmap_status_map = [
('VMA_AREA_NONE', 0 << 0),
('VMA_AREA_REGULAR', 1 << 0),
('VMA_AREA_STACK', 1 << 1),
('VMA_AREA_VSYSCALL', 1 << 2),
('VMA_AREA_VDSO', 1 << 3),
('VMA_AREA_HEAP', 1 << 5),
('VMA_FILE_PRIVATE', 1 << 6),
('VMA_FILE_SHARED', 1 << 7),
('VMA_ANON_SHARED', 1 << 8),
('VMA_ANON_PRIVATE', 1 << 9),
('VMA_AREA_SYSVIPC', 1 << 10),
('VMA_AREA_SOCKET', 1 << 11),
('VMA_AREA_VVAR', 1 << 12),
('VMA_AREA_AIORING', 1 << 13),
('VMA_UNSUPP', 1 << 31),
]
rfile_flags_map = [
('O_WRONLY', 0o1),
('O_RDWR', 0o2),
('O_APPEND', 0o2000),
('O_DIRECT', 0o40000),
('O_LARGEFILE', 0o100000),
]
pmap_flags_map = [
('PE_PARENT', 1 << 0),
('PE_LAZY', 1 << 1),
('PE_PRESENT', 1 << 2),
]
flags_maps = {
'mmap.prot': mmap_prot_map,
'mmap.flags': mmap_flags_map,
'mmap.status': mmap_status_map,
'rfile.flags': rfile_flags_map,
'pmap.flags': pmap_flags_map,
}
gen_maps = {
'task_state': {
1: 'Alive',
3: 'Zombie',
6: 'Stopped'
},
}
sk_maps = {
'family': {
1: 'UNIX',
2: 'INET',
10: 'INET6',
16: 'NETLINK',
17: 'PACKET'
},
'type': {
1: 'STREAM',
2: 'DGRAM',
3: 'RAW',
5: 'SEQPACKET',
10: 'PACKET'
},
'state': {
1: 'ESTABLISHED',
2: 'SYN_SENT',
3: 'SYN_RECV',
4: 'FIN_WAIT1',
5: 'FIN_WAIT2',
6: 'TIME_WAIT',
7: 'CLOSE',
8: 'CLOSE_WAIT',
9: 'LAST_ACK',
10: 'LISTEN'
},
'proto': {
0: 'IP',
6: 'TCP',
17: 'UDP',
136: 'UDPLITE'
},
}
gen_rmaps = {
k: {v2: k2
for k2, v2 in list(v.items())}
for k, v in list(gen_maps.items())
}
sk_rmaps = {
k: {v2: k2
for k2, v2 in list(v.items())}
for k, v in list(sk_maps.items())
}
dict_maps = {
'gen': (gen_maps, gen_rmaps),
'sk': (sk_maps, sk_rmaps),
}
def map_flags(value, flags_map):
bs = [x[0] for x in [x for x in flags_map if value & x[1]]]
value &= ~sum([x[1] for x in flags_map])
if value:
bs.append("0x%x" % value)
return " | ".join(bs)
def unmap_flags(value, flags_map):
if value == '':
return 0
bd = dict(flags_map)
return sum([
int(str(bd.get(x, x)), 0)
for x in [x.strip() for x in value.split('|')]
])
kern_minorbits = 20 # This is how kernel encodes dev_t in new format
def decode_dev(field, value):
if _marked_as_odev(field):
return "%d:%d" % (os.major(value), os.minor(value))
else:
return "%d:%d" % (value >> kern_minorbits,
value & ((1 << kern_minorbits) - 1))
def encode_dev(field, value):
dev = [int(x) for x in value.split(':')]
if _marked_as_odev(field):
return os.makedev(dev[0], dev[1])
else:
return dev[0] << kern_minorbits | dev[1]
def encode_base64(value):
return base64.encodebytes(value)
def decode_base64(value):
return base64.decodebytes(value)
def encode_unix(value):
return quopri.encodestring(value)
def decode_unix(value):
return quopri.decodestring(value)
encode = {'unix_name': encode_unix}
decode = {'unix_name': decode_unix}
def get_bytes_enc(field):
c = _custom_conv(field)
if c:
return encode[c]
else:
return encode_base64
def get_bytes_dec(field):
c = _custom_conv(field)
if c:
return decode[c]
else:
return decode_base64
def is_string(value):
# Python 3 compatibility
if "basestring" in __builtins__:
string_types = basestring # noqa: F821
else:
string_types = (str, bytes)
return isinstance(value, string_types)
def _pb2dict_cast(field, value, pretty=False, is_hex=False):
if not is_hex:
is_hex = _marked_as_hex(field)
if field.type == FD.TYPE_MESSAGE:
return pb2dict(value, pretty, is_hex)
elif field.type == FD.TYPE_BYTES:
return get_bytes_enc(field)(value)
elif field.type == FD.TYPE_ENUM:
return field.enum_type.values_by_number.get(value, None).name
elif field.type in _basic_cast:
cast = _basic_cast[field.type]
if pretty and (cast == int):
if is_hex:
# Fields that have (criu).hex = true option set
# should be stored in hex string format.
return "0x%x" % value
if _marked_as_dev(field):
return decode_dev(field, value)
flags = _marked_as_flags(field)
if flags:
try:
flags_map = flags_maps[flags]
except Exception:
return "0x%x" % value # flags are better seen as hex anyway
else:
return map_flags(value, flags_map)
dct = _marked_as_dict(field)
if dct:
return dict_maps[dct][0][field.name].get(value, cast(value))
return cast(value)
else:
raise Exception("Field(%s) has unsupported type %d" %
(field.name, field.type))
def pb2dict(pb, pretty=False, is_hex=False):
"""
Convert protobuf msg to dictionary.
Takes a protobuf message and returns a dict.
"""
d = collections.OrderedDict() if pretty else {}
for field, value in pb.ListFields():
if field.label == FD.LABEL_REPEATED:
d_val = []
if pretty and _marked_as_ip(field):
if len(value) == 1:
v = socket.ntohl(value[0])
addr = IPv4Address(v)
else:
v = 0 + (socket.ntohl(value[0]) << (32 * 3)) + \
(socket.ntohl(value[1]) << (32 * 2)) + \
(socket.ntohl(value[2]) << (32 * 1)) + \
(socket.ntohl(value[3]))
addr = IPv6Address(v)
d_val.append(addr.compressed)
else:
for v in value:
d_val.append(_pb2dict_cast(field, v, pretty, is_hex))
else:
d_val = _pb2dict_cast(field, value, pretty, is_hex)
d[field.name] = d_val
return d
def _dict2pb_cast(field, value):
# Not considering TYPE_MESSAGE here, as repeated
# and non-repeated messages need special treatment
# in this case, and are hadled separately.
if field.type == FD.TYPE_BYTES:
return get_bytes_dec(field)(value)
elif field.type == FD.TYPE_ENUM:
return field.enum_type.values_by_name.get(value, None).number
elif field.type in _basic_cast:
cast = _basic_cast[field.type]
if (cast == int) and is_string(value):
if _marked_as_dev(field):
return encode_dev(field, value)
flags = _marked_as_flags(field)
if flags:
try:
flags_map = flags_maps[flags]
except Exception:
pass # Try to use plain string cast
else:
return unmap_flags(value, flags_map)
dct = _marked_as_dict(field)
if dct:
ret = dict_maps[dct][1][field.name].get(value, None)
if ret is None:
ret = cast(value, 0)
return ret
# Some int or long fields might be stored as hex
# strings. See _pb2dict_cast.
return cast(value, 0)
else:
return cast(value)
else:
raise Exception("Field(%s) has unsupported type %d" %
(field.name, field.type))
def dict2pb(d, pb):
"""
Convert dictionary to protobuf msg.
Takes dict and protobuf message to be merged into.
"""
for field in pb.DESCRIPTOR.fields:
if field.name not in d:
continue
value = d[field.name]
if field.label == FD.LABEL_REPEATED:
pb_val = getattr(pb, field.name, None)
if is_string(value[0]) and _marked_as_ip(field):
val = ip_address(value[0])
if val.version == 4:
pb_val.append(socket.htonl(int(val)))
elif val.version == 6:
ival = int(val)
pb_val.append(socket.htonl((ival >> (32 * 3)) & 0xFFFFFFFF))
pb_val.append(socket.htonl((ival >> (32 * 2)) & 0xFFFFFFFF))
pb_val.append(socket.htonl((ival >> (32 * 1)) & 0xFFFFFFFF))
pb_val.append(socket.htonl((ival >> (32 * 0)) & 0xFFFFFFFF))
else:
raise Exception("Unknown IP address version %d" %
val.version)
continue
for v in value:
if field.type == FD.TYPE_MESSAGE:
dict2pb(v, pb_val.add())
else:
pb_val.append(_dict2pb_cast(field, v))
else:
if field.type == FD.TYPE_MESSAGE:
# SetInParent method acts just like has_* = true in C,
# and helps to properly treat cases when we have optional
# field with empty repeated inside.
getattr(pb, field.name).SetInParent()
dict2pb(value, getattr(pb, field.name, None))
else:
setattr(pb, field.name, _dict2pb_cast(field, value))
return pb