diff --git a/libs/sm/blktap2.py b/libs/sm/blktap2.py index edea11a22..75fd44544 100644 --- a/libs/sm/blktap2.py +++ b/libs/sm/blktap2.py @@ -29,6 +29,7 @@ import syslog as _syslog import glob import json +import base64 from syslog import openlog, syslog from stat import * # S_ISBLK(), ... @@ -379,8 +380,7 @@ def open(cls, pid, minor, _type, _file, options): if not key: raise util.SMException("No key found with key hash {}".format(key_hash)) - input = key - text_mode = False + input = base64.b64encode(key).decode('utf-8') args.append('-E') cls._pread(args=args, input=input, text_mode=text_mode) diff --git a/libs/sm/cbtutil.py b/libs/sm/cbtutil.py index 8b0aa4c75..f4dceef80 100644 --- a/libs/sm/cbtutil.py +++ b/libs/sm/cbtutil.py @@ -19,6 +19,7 @@ from sm.core import util import uuid +import base64 from sm.constants import CBT_UTIL @@ -81,10 +82,10 @@ def get_cbt_consistency(file_name): def get_cbt_bitmap(file_name): """Get bitmap field from log file""" - cmd = [CBT_UTIL, "get", "-n", file_name, "-b"] - ret = _call_cbt_util(cmd, text=False) - # Do not strip the return string. It's a byte string and stripping - # it sometimes leads to loss of information + cmd = [CBT_UTIL, "get", "-n", file_name, "-b", "-E"] + ret = _call_cbt_util(cmd) + if ret is not None: + ret = base64.b64decode(ret.strip()) return ret diff --git a/libs/sm/vhdutil.py b/libs/sm/vhdutil.py index ccb5e57b4..1f6eb3253 100644 --- a/libs/sm/vhdutil.py +++ b/libs/sm/vhdutil.py @@ -21,6 +21,7 @@ import errno import zlib import re +import base64 from sm.core import xs_errors import time @@ -298,8 +299,10 @@ def getDepth(path): def getBlockBitmap(path): - cmd = [VHD_UTIL, "read", OPT_LOG_ERR, "-B", "-n", path] - text = ioretry(cmd, text=False) + cmd = [VHD_UTIL, "read", OPT_LOG_ERR, "-B", "-n", path, "-E"] + text = ioretry(cmd) + if text is not None: + text = base64.b64decode(text.strip()) return zlib.compress(text) diff --git a/tests/test_cbt.py b/tests/test_cbt.py index 865f74f2e..3121602c7 100644 --- a/tests/test_cbt.py +++ b/tests/test_cbt.py @@ -764,7 +764,8 @@ def test_list_changed_blocks_strip_sensitive_bitmap(self, context, mock_lock, data = b'\x0c\xbb3a\xf75Jy\x17\x04\x92;\x0b\x93\xf3E' bitmap2.frombytes(data) mock_size.side_effect = [vdi_size1, vdi_size2] - mock_call.side_effect = [bitmap1.tobytes(), data] + mock_call.side_effect = [base64.b64encode(bitmap1.tobytes()).decode(), + base64.b64encode(data).decode()] bitmap1.bytereverse() bitmap2.bytereverse() expected_string = base64.b64encode((bitmap1 | bitmap2).tobytes()).decode() diff --git a/tests/test_vhdutil.py b/tests/test_vhdutil.py index 0af601397..245316a75 100644 --- a/tests/test_vhdutil.py +++ b/tests/test_vhdutil.py @@ -2,6 +2,7 @@ import unittest import unittest.mock as mock import zlib +import base64 from sm import lvhdutil from sm import vhdutil @@ -147,8 +148,8 @@ def test_get_block_bitmap(self, context): def test_function(args, inp): nonlocal call_args call_args = args - # Not a real bitmap data - return 0, b"some dummy text", b"" + # Since vhd-util now returns base64 encoded data with -E flag + return 0, base64.b64encode(b"some dummy text"), b"" context.add_executable(VHD_UTIL, test_function) @@ -160,7 +161,7 @@ def test_function(args, inp): self.assertEqual("some dummy text", zlib.decompress(result).decode()) self.assertEqual([ VHD_UTIL, "read", "--debug", "-B", - "-n", TEST_VHD_NAME], + "-n", TEST_VHD_NAME, "-E"], call_args) @testlib.with_context