Mirrors/criu
1
0
Fork 0
mirror of https://github.com/checkpoint-restore/criu.git synced 2026-07-17 16:47:50 +00:00
Code Issues Projects Releases Wiki Activity
criu/lib/py/images/pb2dict.py
Andrei Vagin 5aa72e7237 py: Reformat everything into pep8 style 
As discussed on the mailing list, current .py files formatting does not
conform to the world standard, so we should better reformat it. For this
the yapf tool is used. The command I used was

  yapf -i $(find -name *.py)

Signed-off-by: Pavel Emelyanov <xemul@virtuozzo.com>
2019-09-07 15:59:56 +03:00

424 lines
12 KiB
Python
Raw Blame History

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
Reference in a new issue View git blame Copy permalink