def test_key_material_read(self): volume = Volume(disk=Disk(ImageParser(keys={'3': 'hello'}), "..."), index='3') self.assertEqual(volume.key, "hello") volume = Volume(disk=Disk(ImageParser(keys={'3': 'hello', '*': 'ola'}), "..."), index='2') self.assertEqual(volume.key, "ola") volume = Volume(disk=Disk(ImageParser(keys={'3': 'hello'}), "..."), index='1') self.assertEqual(volume.key, "")
def preloop(self): """if the parser is not already set, loads the parser.""" if not self.parser: self.stdout.write("Welcome to imagemounter {version}".format(version=__version__)) self.stdout.write("\n") self.parser = ImageParser() for p in self.args.paths: self.onecmd('disk "{}"'.format(p))
def test_no_root(self): parser = ImageParser() disk = parser.add_disk("...") v1 = Volume(disk) v1.mountpoint = '...' v1.info['lastmountpoint'] = '/etc/x' v2 = Volume(disk) v2.mountpoint = '....' v2.info['lastmountpoint'] = '/etc' disk.volumes.volumes = [v1, v2] with self.assertRaises(NoRootFoundError): parser.reconstruct()
def test_simple(self): parser = ImageParser() disk = parser.add_disk("...") v1 = Volume(disk) v1.mountpoint = '...' v1.info['lastmountpoint'] = '/' v2 = Volume(disk) v2.mountpoint = '....' v2.info['lastmountpoint'] = '/etc' disk.volumes.volumes = [v1, v2] with mock.patch.object(v2, "bindmount") as v2_bm: parser.reconstruct() v2_bm.assert_called_once_with(".../etc")
def test_combination(self): # Add values here to test full combinations of specific filesystem types # The _ as key is the expected result _ = "use as key for the expected result" definitions = [ {_: "cramfs", "blkid": "cramfs", "magic": "Linux Compressed ROM File System data", "fsdescription": "???"}, {_: "exfat", "blkid": "exfat", "fsdescription": "NTFS / exFAT", "statfstype": "exFAT"}, {_: "ext", "blkid": "ext4", "fsdescription": "Linux (0x83)", "guid": "", "statfstype": "Ext2"}, {_: "fat", "blkid": "vfat", "magic": "FAT (12 bit)", "fsdescription": "DOS FAT12 (0x04)", "statfstype": "FAT12"}, {_: "iso", "blkid": "iso9660", "magic": ".. ISO 9660 ..", "statfstype": "ISO9660"}, {_: "minix", "blkid": "min ix", "magic": "Minix filesystem", "fsdescription": "???"}, {_: "ntfs", "blkid": "ntfs", "magic": ".. NTFS ..", "fsdescription": "NTFS / exFAT", "statfstype": "NTFS"}, {_: "squashfs", "blkid": "squashfs", "magic": "Squashfs filesystem", "fsdescription": "???"}, {_: "lvm", "guid": "79D3D6E6-07F5-C244-A23C-238F2A3DF928"}, {_: "raid", "fsdescription": "Linux (0x83)", "blkid": "linux_raid_member"}, {_: "volumesystem", "blkid": "dos", "fsdescription": "Logical Volume"}, {_: "volumesystem", "blkid": "dos", "fsdescription": "RAID Volume"}, {_: "volumesystem", "blkid": "dos", "magic": "DOS/MBR boot sector"}, {_: "volumesystem", "fsdescription": "BSD/386, 386BSD, NetBSD, FreeBSD (0xa5)", "blkid": "ufs"}, {_: "ufs", "fsdescription": "4.2BSD (0x07)", "blkid": "ufs"}, ] for definition in definitions: volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=definition.get("blkid")) volume._get_magic_type = mock.Mock(return_value=definition.get("magic")) volume.info = definition volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES[definition[_]], volume.fstype)
def test_fstype_fallback(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume.fstype = "?bsd" volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.determine_fs_type() self.assertEqual("bsd", volume.fstype)
def test_magic(self): # Add values here that are shown in the wild for file magic output # !! Always try to add it also to test_combination descriptions = { "Linux Compressed ROM File System data": "cramfs", "Linux rev 1.0 ext2 filesystem data": "ext", 'DOS/MBR boot sector, code offset 0x3c+2, OEM-ID "mkfs.fat", sectors/cluster 4, ' 'root entries 512, sectors 100 (volumes <=32 MB) , Media descriptor 0xf8, ' 'sectors/FAT 1, sectors/track 32, heads 64, serial number 0x3cb7474b, ' 'label: "TEST ", FAT (12 bit)': "fat", "ISO 9660 CD-ROM filesystem data 'test'": "iso", "Minix filesystem, V1, 30 char names, 12800 zones": "minix", 'DOS/MBR boot sector, code offset 0x52+2, OEM-ID "NTFS ", sectors/cluster 8, ' 'Media descriptor 0xf8, sectors/track 0, dos < 4.0 BootSector (0x80), FAT (1Y bit ' 'by descriptor); NTFS, sectors 2048, $MFT start cluster 4, $MFTMirror start ' 'cluster 128, bytes/RecordSegment 2^(-1*246), clusters/index block 1, serial ' 'number 04e8742c12a96cecd; contains Microsoft Windows XP/VISTA bootloader BOOTMGR': "ntfs", "Squashfs filesystem, little endian, version 4.0": "squashfs", } for description, fstype in descriptions.items(): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=description) volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES[fstype], volume.fstype)
def test_fstype_fallback(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume._get_fstype_from_parser('?ufs') volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES["ufs"], volume.fstype)
def test_ext4(self): # Removed some items from this output as we don't use it in its entirety anyway result = b"""FILE SYSTEM INFORMATION -------------------------------------------- File System Type: Ext4 Volume Name: Example Volume ID: 2697f5b0479b15b1b4c81994387cdba Last Written at: 2017-07-02 12:23:22 (CEST) Last Checked at: 2016-07-09 20:27:28 (CEST) Last Mounted at: 2017-07-02 12:23:23 (CEST) Unmounted properly Last mounted on: / Source OS: Linux BLOCK GROUP INFORMATION --------------------------------------------""" with mock.patch('subprocess.Popen') as mock_popen: type(mock_popen()).stdout = mock.PropertyMock(return_value=io.BytesIO(result)) volume = Volume(disk=Disk(ImageParser(), "...")) volume.get_raw_path = mock.Mock(return_value="...") volume._load_fsstat_data() self.assertEqual(volume.info['statfstype'], 'Ext4') self.assertEqual(volume.info['lastmountpoint'], '/') self.assertEqual(volume.info['label'], '/ (Example)') self.assertEqual(volume.info['version'], 'Linux') # must be called after reading BLOCK GROUP INFORMATION mock_popen().terminate.assert_called()
def test_luks_key_communication(self, _, check_call): def modified_check_call(cmd, *args, **kwargs): if cmd[0:2] == ['cryptsetup', 'isLuks']: return True return mock.DEFAULT check_call.side_effect = modified_check_call original_popen = subprocess.Popen def modified_popen(cmd, *args, **kwargs): if cmd[0:3] == ['cryptsetup', '-r', 'luksOpen']: # A command that requests user input x = original_popen(["python2", "-c", "print(raw_input(''))"], *args, **kwargs) return x return mock.DEFAULT with mock.patch("subprocess.Popen", side_effect=modified_popen) as popen: disk = Disk(ImageParser(keys={'1': 'p:passphrase'}), "...") disk.is_mounted = True volume = Volume(disk=disk, fstype='luks', index='1', parent=disk) volume.mount() self.assertTrue(volume.is_mounted) self.assertEqual(len(volume.volumes), 1) self.assertEqual(volume.volumes[0].info['fsdescription'], "LUKS Volume")
def test_fsdescription(self): # Add names in here that are shown in the wild for output of mmls / gparted # !! Always try to add it also to test_combination descriptions = { # Names assigned by imagemounter "Logical Volume": "unknown", "LUKS Volume": "unknown", "BDE Volume": "unknown", "RAID Volume": "unknown", "VSS Store": "unknown", "NTFS / exFAT": "ntfs", # mmls, should use fallback "Linux (0x83)": "unknown", # should use unknown "4.2BSD": "ufs", "BSD/386, 386BSD, NetBSD, FreeBSD (0xa5)": "volumesystem", "DOS FAT16": "fat", } for description, fstype in descriptions.items(): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) self.fstype = "" # prevent fallback to unknown by default volume.info['fsdescription'] = description volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES[fstype], volume.fstype)
def test_guid(self): for description, fstype in FILE_SYSTEM_GUIDS.items(): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.info['guid'] = description volume.determine_fs_type() self.assertEqual(fstype, volume.fstype)
def test_multiple_roots(self): parser = ImageParser() disk = parser.add_disk("...") v1 = Volume(disk) v1.index = '1' v1.mountpoint = '...' v1.info['lastmountpoint'] = '/' v2 = Volume(disk) v2.index = '2' v2.mountpoint = '....' v2.info['lastmountpoint'] = '/' v3 = Volume(disk) v3.index = '3' v3.mountpoint = '.....' v3.info['lastmountpoint'] = '/etc' disk.volumes.volumes = [v1, v2, v3] with mock.patch.object(v1, "bindmount") as v1_bm, mock.patch.object(v2, "bindmount") as v2_bm, \ mock.patch.object(v3, "bindmount") as v3_bm: parser.reconstruct() v1_bm.assert_not_called() v2_bm.assert_not_called() v3_bm.assert_called_with('.../etc')
def test_parted_requests_input(self, check_output): def modified_command(cmd, *args, **kwargs): if cmd[0] == 'parted': # A command that requests user input return check_output_(["python2", "-c", "exec \"try: raw_input('>> ')\\nexcept: pass\""], *args, **kwargs) return mock.DEFAULT check_output.side_effect = modified_command disk = Disk(ImageParser(), path="...") list(disk.volumes.detect_volumes(method='parted')) check_output.assert_called()
def test_killed_after_timeout(self): def mock_side_effect(*args, **kwargs): time.sleep(0.2) return io.BytesIO(b"") with mock.patch('subprocess.Popen') as mock_popen: type(mock_popen()).stdout = mock.PropertyMock(side_effect=mock_side_effect) volume = Volume(disk=Disk(ImageParser(), "...")) volume.get_raw_path = mock.Mock(return_value="...") volume._load_fsstat_data(timeout=0.1) mock_popen().terminate.assert_called()
def test_fstype_fallback_unknown(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.info['fsdescription'] = "Linux (0x83)" # If something more specific is set, we use that volume.fstype = "?bsd" volume.determine_fs_type() self.assertEqual("bsd", volume.fstype) # Otherwise we fallback to unknown if Linux (0x83) is set volume.fstype = "" volume.determine_fs_type() self.assertEqual("unknown", volume.fstype)
def test_fstype_fallback_unknown(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.info['fsdescription'] = "Linux (0x83)" # If something more specific is set, we use that volume._get_fstype_from_parser('?ufs') volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES["ufs"], volume.fstype) # Otherwise we fallback to unknown if Linux (0x83) is set volume._get_fstype_from_parser('') volume.determine_fs_type() self.assertEqual(UnknownFileSystemType(), volume.fstype)
def test_utf8_label(self): # Removed some items from this output as we don't use it in its entirety anyway result = b"""FILE SYSTEM INFORMATION -------------------------------------------- File System Type: Ext4 Volume Name: \xd0\xa0\xd0\xbe\xd1\x81\xd1\x81\xd0\xb8\xd0\xb8 Volume ID: 2697f5b0479b15b1b4c81994387cdba""" with mock.patch('subprocess.Popen') as mock_popen: type(mock_popen()).stdout = mock.PropertyMock(return_value=io.BytesIO(result)) volume = Volume(disk=Disk(ImageParser(), "...")) volume.get_raw_path = mock.Mock(return_value="...") volume._load_fsstat_data() self.assertEqual(volume.info['statfstype'], 'Ext4') self.assertEqual(volume.info['label'], u'\u0420\u043e\u0441\u0441\u0438\u0438')
def test_ntfs(self): # Removed some items from this output as we don't use it in its entirety anyway result = b"""FILE SYSTEM INFORMATION -------------------------------------------- File System Type: NTFS Volume Serial Number: 4E8742C12A96CECD OEM Name: NTFS Version: Windows XP""" with mock.patch('subprocess.Popen') as mock_popen: type(mock_popen()).stdout = mock.PropertyMock(return_value=io.BytesIO(result)) volume = Volume(disk=Disk(ImageParser(), "...")) volume.get_raw_path = mock.Mock(return_value="...") volume._load_fsstat_data() self.assertEqual(volume.info['statfstype'], 'NTFS') self.assertNotIn("lastmountpoint", volume.info) self.assertNotIn("label", volume.info) self.assertEqual(volume.info['version'], 'Windows XP')
def test_guid(self): FILE_SYSTEM_GUIDS = { '2AE031AA-0F40-DB11-9590-000C2911D1B8': 'vmfs', # '8053279D-AD40-DB11-BF97-000C2911D1B8': 'vmkcore-diagnostics', # '6A898CC3-1DD2-11B2-99A6-080020736631': 'zfs-member', # 'C38C896A-D21D-B211-99A6-080020736631': 'zfs-member', # '0FC63DAF-8483-4772-8E79-3D69D8477DE4': 'linux', 'E6D6D379-F507-44C2-A23C-238F2A3DF928': 'lvm', '79D3D6E6-07F5-C244-A23C-238F2A3DF928': 'lvm', 'CA7D7CCB-63ED-4C53-861C-1742536059CC': 'luks' } for description, fstype in FILE_SYSTEM_GUIDS.items(): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.info['guid'] = description volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES[fstype], volume.fstype)
def test_blkid(self): # Add values here that are shown in the wild for blkid # !! Always try to add it also to test_combination descriptions = { "cramfs": "cramfs", "ext4": "ext", "ext2": "ext", "vfat": "fat", "iso9660": "iso", "minix": "minix", "ntfs": "ntfs", "squashfs": "squashfs", "ufs": "ufs", "dos": "volumesystem", } for description, fstype in descriptions.items(): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=description) volume._get_magic_type = mock.Mock(return_value=None) volume.determine_fs_type() self.assertEqual(fstype, volume.fstype)
def test_no_clue_fstype(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value=None) volume._get_magic_type = mock.Mock(return_value=None) volume.determine_fs_type() self.assertEqual("unknown", volume.fstype)
def preloop(self): self.stdout.write("Welcome to imagemounter {version}".format(version=__version__)) self.stdout.write("\n") self.parser = ImageParser()
class ImageMounterShell(ArgumentParsedShell): prompt = '(imount) ' file = None parser = None def preloop(self): self.stdout.write("Welcome to imagemounter {version}".format(version=__version__)) self.stdout.write("\n") self.parser = ImageParser() def error(self, error): self.stdout.write('*** %s\n' % error) def parser_disk(self, parser): parser.description = "Add a disk to the current parser" p = parser.add_argument('path', help='path to the disk image that you want to mount') try: from argcomplete.completers import FilesCompleter p.completer = FilesCompleter([".dd", ".e01", ".aff", ".DD", ".E01", ".AFF"]) except ImportError: pass parser.add_argument("--mounter", choices=DISK_MOUNTERS, help="the method to mount with") def arg_disk(self, args): if not os.path.exists(args.path): return self.error("The path {path} does not exist".format(path=args.path)) disk = self.parser.add_disk(args.path, disk_mounter=args.mounter) disk.mount() for _ in disk.detect_volumes(): pass print("Added {path} to the image mounter as index {index}".format(path=args.path, index=disk.index)) def _get_all_indexes(self): if self.parser: return [v.index for v in self.parser.get_volumes()] + [d.index for d in self.parser.disks] else: return None def parser_mount(self, parser): parser.description = "Mount a volume by its index" parser.add_argument('index', help='volume index', choices=self._get_all_indexes()) parser.add_argument('-r', '--recursive', action='store_true') def arg_mount(self, args): col = get_coloring_func() volume_or_disk = self.parser.get_by_index(args.index) volume, disk = (volume_or_disk, None) if not isinstance(volume_or_disk, Disk) else (None, volume_or_disk) if not args.recursive: if disk and not disk.is_mounted: try: disk.mount() except Exception as e: pass else: try: volume.init_volume() if volume.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(volume.mountpoint, "green", attrs=['bold']), index=volume.index)) else: print("Mounted volume {index} (no mountpoint available)".format(index=volume.index)) except Exception as e: import traceback ; traceback.print_exc() print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=volume.index), "red")) else: if disk: it = disk.init_volumes else: it = volume.init for v in it(): if v.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(v.mountpoint, "green", attrs=['bold']), index=v.index)) elif v.exception: e = v.exception print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=v.index), "red")) def parser_unmount(self, parser): parser.description = "Unmount a volume by its index" parser.add_argument('index', help='volume index', nargs='?', choices=self._get_all_indexes()) def arg_unmount(self, args): if args.index: volume = self.parser.get_by_index(args.index) volume.unmount() print("Unmounted {index}".format(index=volume.index)) else: self.parser.clean() print("Unmounted everything") def do_show(self, args): col = get_coloring_func() for disk in self.parser.disks: print("- {index:<5} {type} {filename}" .format(index=col("{:<5}".format(disk.index), 'green' if disk.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(disk.volumes.vstype), attrs=['dark']), filename=disk.paths[0])) def _show_volume_system(volumes, level=0): level += 1 for i, v in enumerate(volumes): level_str = " "*level + ("└ " if i == len(volumes)-1 else "├ ") tp = v.volumes.vstype if v.fstype == 'volumesystem' else v.fstype if v.flag == 'alloc' else v.flag print("{level_str}{index} {type} {size:<10} {description}" .format(level_str=level_str, index=col("{:<5}".format(v.index), 'green' if v.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(tp), attrs=['dark']), description=v.get_description(with_index=False, with_size=False)[:30], size=v.get_formatted_size())) _show_volume_system(v.volumes, level) _show_volume_system(disk.volumes) def do_quit(self, arg): """Quits the program.""" return True
def test_no_volumes(self): parser = ImageParser() parser.add_disk("...") with self.assertRaises(NoRootFoundError): parser.reconstruct()
def test_valid_vstype(self): volume = Volume(disk=Disk(ImageParser(), "..."), fstype="dos") volume.determine_fs_type() self.assertEqual("dos", volume.volumes.vstype) self.assertEqual(FILE_SYSTEM_TYPES["volumesystem"], volume.fstype)
class ImageMounterShell(ArgumentParsedShell): prompt = '(imount) ' file = None parser = None args = None saved = False def save(self): with open(SAVE_PATH, 'wb') as f: pickle.dump(self.parser, f) self.saved = True def load(self): with open(SAVE_PATH, 'rb') as f: self.parser = pickle.load(f) def preloop(self): """if the parser is not already set, loads the parser.""" if not self.parser: self.stdout.write("Welcome to imagemounter {version}".format(version=__version__)) self.stdout.write("\n") self.parser = ImageParser() for p in self.args.paths: self.onecmd('disk "{}"'.format(p)) def onecmd(self, line): """Do not crash the entire program when a single command fails.""" try: return cmd.Cmd.onecmd(self, line) except Exception as e: print("Critical error.", e) def error(self, error): """Writes an error to the console""" self.stdout.write('*** %s\n' % error) def _get_all_indexes(self): """Returns all indexes available in the parser""" if self.parser: return [v.index for v in self.parser.get_volumes()] + [d.index for d in self.parser.disks] else: return None def _get_by_index(self, index): """Returns a volume,disk tuple for the specified index""" volume_or_disk = self.parser.get_by_index(index) volume, disk = (volume_or_disk, None) if not isinstance(volume_or_disk, Disk) else (None, volume_or_disk) return volume, disk ###################################################################### # disk command ###################################################################### def parser_disk(self, parser): parser.description = "Add a disk to the current parser" p = parser.add_argument('path', help='path to the disk image that you want to mount') try: from argcomplete.completers import FilesCompleter p.completer = FilesCompleter([".dd", ".e01", ".aff", ".DD", ".E01", ".AFF"]) except ImportError: pass parser.add_argument("--mounter", choices=DISK_MOUNTERS, help="the method to mount with") def arg_disk(self, args): args.path = os.path.expanduser(args.path) if not os.path.exists(args.path): return self.error("The path {path} does not exist".format(path=args.path)) disk = self.parser.add_disk(args.path, disk_mounter=args.mounter) disk.mount() for _ in disk.detect_volumes(): pass print("Added {path} to the image mounter as index {index}".format(path=args.path, index=disk.index)) ###################################################################### # mount command ###################################################################### def parser_mount(self, parser): parser.description = "Mount a volume or disk by its index" parser.add_argument('index', help='volume or disk index', choices=self._get_all_indexes()) parser.add_argument('-r', '--recursive', action='store_true', help='recursively mount all volumes under this index') parser.add_argument('-f', '--fstype', default=None, choices=FILE_SYSTEM_TYPES, help='specify the file system type for the volume') parser.add_argument('-k', '--key', default=None, help='specify the key for the volume') def arg_mount(self, args): col = get_coloring_func() volume, disk = self._get_by_index(args.index) if not args.recursive: if disk: if not disk.is_mounted: try: disk.mount() except Exception as e: pass else: print(col("Disk {} is already mounted.".format(disk.index), 'red')) else: if not volume.is_mounted: try: if args.key is not None: volume.key = args.key volume.init_volume(fstype=args.fstype) if volume.is_mounted: if volume.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(volume.mountpoint, "green", attrs=['bold']), index=volume.index)) else: print("Mounted volume {index} (no mountpoint available)".format(index=volume.index)) else: print("Refused to mount volume {index}.".format(index=volume.index)) except Exception as e: print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=volume.index), "red")) else: if volume.mountpoint: print(col("Volume {} is already mounted at {}.".format(volume.index, volume.mountpoint), 'red')) else: print(col("Volume {} is already mounted.".format(volume.index), 'red')) else: if disk: if not disk.is_mounted: try: disk.mount() except Exception as e: pass it = disk.init_volumes else: it = volume.init for v in it(): if v.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(v.mountpoint, "green", attrs=['bold']), index=v.index)) elif v.exception: e = v.exception print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=v.index), "red")) ###################################################################### # unmount command ###################################################################### def parser_unmount(self, parser): parser.description = "Unmount a disk or volume by its index. Is recursive by design." parser.add_argument('index', help='volume index', nargs='?', choices=self._get_all_indexes()) def arg_unmount(self, args): if args.index: volume = self.parser.get_by_index(args.index) volume.unmount() print("Unmounted {index}".format(index=volume.index)) else: self.parser.clean() print("Unmounted everything") ###################################################################### # show command ###################################################################### def parser_show(self, parser): parser.description = "Without arguments, displays a tree view of all known volumes and disks. With an " \ "argument, it provides details about the referenced volume or disk." parser.add_argument('index', help='volume index', nargs='?', choices=self._get_all_indexes()) def arg_show(self, args): if not args.index: self._show_tree() else: col = get_coloring_func() volume, disk = self._get_by_index(args.index) # displays a volume's details if volume: print(col("Details of volume {description}".format(description=volume.get_description()), attrs=['bold'])) if volume.is_mounted: print(col("Currently mounted", 'green')) print("\nAttributes:") for attr in ['size', 'offset', 'slot', 'flag', 'block_size', 'fstype', 'key']: print(" - {name:<15} {value}".format(name=attr, value=getattr(volume, attr))) if volume.mountpoint or volume.loopback or volume._paths: print("\nPaths:") if volume.mountpoint: print(" - mountpoint {value}".format(value=volume.mountpoint)) if volume.loopback: print(" - loopback {value}".format(value=volume.loopback)) for k, v in volume._paths.items(): print(" - {name:<15} {value}".format(name=k, value=v)) if volume.info: print("\nAdditional information:") for k, v in volume.info.items(): print(" - {name:<15} {value}".format(name=k, value=v)) else: # displays a disk's details print(col("Details of disk {description}".format(description=disk.paths[0]), attrs=['bold'])) if disk.is_mounted: print(col("Currently mounted", 'green')) print("\nAttributes:") for attr in ['offset', 'block_size', 'read_write', 'disk_mounter']: print(" - {name:<15} {value}".format(name=attr, value=getattr(disk, attr))) if disk.mountpoint or disk.rwpath or disk._paths: print("\nPaths:") if disk.mountpoint: print(" - mountpoint {value}".format(value=disk.mountpoint)) if disk.rwpath: print(" - rwpath {value}".format(value=disk.rwpath)) for k, v in disk._paths.items(): print(" - {name:<15} {value}".format(name=k, value=v)) def _show_tree(self): col = get_coloring_func() for disk in self.parser.disks: print("- {index:<5} {type} {filename}".format( index=col("{:<5}".format(disk.index), 'green' if disk.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(disk.volumes.vstype), attrs=['dark']), filename=disk.paths[0] )) def _show_volume_system(volumes, level=0): level += 1 for i, v in enumerate(volumes): level_str = " " * level + ("└ " if i == len(volumes) - 1 else "├ ") tp = v.volumes.vstype if v.fstype.type == 'volumesystem' else v.fstype if v.flag == 'alloc' else v.flag print("{level_str}{index} {type} {size:<10} {description}".format( level_str=level_str, index=col("{:<5}".format(v.index), 'green' if v.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(str(tp)), attrs=['dark']), description=v.get_description(with_index=False, with_size=False)[:30], size=v.get_formatted_size() )) _show_volume_system(v.volumes, level) _show_volume_system(disk.volumes) ###################################################################### # set command ###################################################################### def parser_set(self, parser): parser.description = "Modifies a property of a volume or disk. This is for advanced usage only." parser.add_argument('index', help='volume index', choices=self._get_all_indexes()) parser.add_argument('name', help='property name', choices=['size', 'offset', 'slot', 'flag', 'block_size', 'fstype', 'key', 'disk_mounter']) parser.add_argument('value', help='property value') def arg_set(self, args): col = get_coloring_func() volume, disk = self._get_by_index(args.index) if volume: if args.name in ['size', 'offset', 'block_size']: try: setattr(volume, args.name, int(args.value)) except ValueError: print(col("Invalid value provided for {}".format(args.name), 'red')) else: print(col("Updated value for {}".format(args.name), 'green')) elif args.name in ['slot', 'flag', 'fstype', 'key']: setattr(volume, args.name, args.value) print(col("Updated value for {}".format(args.name), 'green')) else: print(col("Property {} can't be set for a volume".format(args.name), 'red')) else: if args.name in ['offset', 'block_size']: try: setattr(disk, args.name, int(args.value)) except ValueError: print(col("Invalid value provided for {}".format(args.name), 'red')) else: print(col("Updated value for {}".format(args.name), 'green')) elif args.name in ['disk_mounter']: setattr(disk, args.name, args.value) print(col("Updated value for {}".format(args.name), 'green')) else: print(col("Property {} can't be set for a disk".format(args.name), 'red')) ###################################################################### # quit command ###################################################################### def do_save(self, arg): self.save() def do_EOF(self, arg): self.save() return True def do_quit(self, arg): """Quits the program.""" if self.saved: self.save() else: self.parser.clean() return True
def test_valid_fstype(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume.fstype = 'ext' volume.determine_fs_type() self.assertEqual("ext", volume.fstype)
def test_valid_fstype(self): volume = Volume(disk=Disk(ImageParser(), "..."), fstype='ext') volume.determine_fs_type() self.assertEqual(FILE_SYSTEM_TYPES['ext'], volume.fstype)
def test_valid_vstype(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume.fstype = 'dos' volume.determine_fs_type() self.assertEqual("dos", volume.volumes.vstype) self.assertEqual("volumesystem", volume.fstype)
def test_read_bytes_crash(self, mock_open): mock_open().__enter__().read.side_effect = IOError volume = Volume(disk=Disk(ImageParser(), "...")) volume.get_raw_path = mock.Mock(return_value="...") self.assertIsNone(volume._get_magic_type())
def test_little_clue_fstype(self): volume = Volume(disk=Disk(ImageParser(), "...")) volume._get_blkid_type = mock.Mock(return_value="-") volume._get_magic_type = mock.Mock(return_value="-") volume.determine_fs_type() self.assertEqual(UnknownFileSystemType(), volume.fstype)
class ImageMounterShell(ArgumentParsedShell): prompt = '(imount) ' file = None parser = None args = None saved = False def save(self): with open(SAVE_PATH, 'wb') as f: pickle.dump(self.parser, f) self.saved = True def load(self): with open(SAVE_PATH, 'rb') as f: self.parser = pickle.load(f) def preloop(self): """if the parser is not already set, loads the parser.""" if not self.parser: self.stdout.write("Welcome to imagemounter {version}".format(version=__version__)) self.stdout.write("\n") self.parser = ImageParser() for p in self.args.paths: self.onecmd('disk "{}"'.format(p)) def onecmd(self, line): """Do not crash the entire program when a single command fails.""" try: return cmd.Cmd.onecmd(self, line) except Exception as e: print("Critical error.", e) def error(self, error): """Writes an error to the console""" self.stdout.write('*** %s\n' % error) def _get_all_indexes(self): """Returns all indexes available in the parser""" if self.parser: return [v.index for v in self.parser.get_volumes()] + [d.index for d in self.parser.disks] else: return None def _get_by_index(self, index): """Returns a volume,disk tuple for the specified index""" volume_or_disk = self.parser.get_by_index(index) volume, disk = (volume_or_disk, None) if not isinstance(volume_or_disk, Disk) else (None, volume_or_disk) return volume, disk ###################################################################### # disk command ###################################################################### def parser_disk(self, parser): parser.description = "Add a disk to the current parser" p = parser.add_argument('path', help='path to the disk image that you want to mount') try: from argcomplete.completers import FilesCompleter p.completer = FilesCompleter([".dd", ".e01", ".aff", ".DD", ".E01", ".AFF"]) except ImportError: pass parser.add_argument("--mounter", choices=DISK_MOUNTERS, help="the method to mount with") def arg_disk(self, args): args.path = os.path.expanduser(args.path) if not os.path.exists(args.path): return self.error("The path {path} does not exist".format(path=args.path)) disk = self.parser.add_disk(args.path, disk_mounter=args.mounter) disk.mount() for _ in disk.detect_volumes(): pass print("Added {path} to the image mounter as index {index}".format(path=args.path, index=disk.index)) ###################################################################### # mount command ###################################################################### def parser_mount(self, parser): parser.description = "Mount a volume or disk by its index" parser.add_argument('index', help='volume or disk index', choices=self._get_all_indexes()) parser.add_argument('-r', '--recursive', action='store_true', help='recursively mount all volumes under this index') parser.add_argument('-f', '--fstype', default=None, choices=FILE_SYSTEM_TYPES, help='specify the file system type for the volume') parser.add_argument('-k', '--key', default=None, help='specify the key for the volume') def arg_mount(self, args): col = get_coloring_func() volume, disk = self._get_by_index(args.index) if not args.recursive: if disk: if not disk.is_mounted: try: disk.mount() except Exception as e: pass else: print(col("Disk {} is already mounted.".format(disk.index), 'red')) else: if not volume.is_mounted: try: if args.key is not None: volume.key = args.key volume.init_volume(fstype=args.fstype) if volume.is_mounted: if volume.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(volume.mountpoint, "green", attrs=['bold']), index=volume.index)) else: print("Mounted volume {index} (no mountpoint available)".format(index=volume.index)) else: print("Refused to mount volume {index}.".format(index=volume.index)) except Exception as e: print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=volume.index), "red")) else: if volume.mountpoint: print(col("Volume {} is already mounted at {}.".format(volume.index, volume.mountpoint), 'red')) else: print(col("Volume {} is already mounted.".format(volume.index), 'red')) else: if disk: if not disk.is_mounted: try: disk.mount() except Exception as e: pass it = disk.init_volumes else: it = volume.init for v in it(): if v.mountpoint: print("Mounted volume {index} at {path}" .format(path=col(v.mountpoint, "green", attrs=['bold']), index=v.index)) elif v.exception: e = v.exception print(col("An error occurred while mounting volume {index}: {type}: {args}" .format(type=type(e).__name__, args=" ".join(map(str, e.args)), index=v.index), "red")) ###################################################################### # unmount command ###################################################################### def parser_unmount(self, parser): parser.description = "Unmount a disk or volume by its index. Is recursive by design." parser.add_argument('index', help='volume index', nargs='?', choices=self._get_all_indexes()) def arg_unmount(self, args): if args.index: volume = self.parser.get_by_index(args.index) volume.unmount() print("Unmounted {index}".format(index=volume.index)) else: self.parser.clean() print("Unmounted everything") ###################################################################### # show command ###################################################################### def parser_show(self, parser): parser.description = "Without arguments, displays a tree view of all known volumes and disks. With an " \ "argument, it provides details about the referenced volume or disk." parser.add_argument('index', help='volume index', nargs='?', choices=self._get_all_indexes()) def arg_show(self, args): if not args.index: self._show_tree() else: col = get_coloring_func() volume, disk = self._get_by_index(args.index) # displays a volume's details if volume: print(col("Details of volume {description}".format(description=volume.get_description()), attrs=['bold'])) if volume.is_mounted: print(col("Currently mounted", 'green')) print("\nAttributes:") for attr in ['size', 'offset', 'slot', 'flag', 'block_size', 'fstype', 'key']: print(" - {name:<15} {value}".format(name=attr, value=getattr(volume, attr))) if volume.mountpoint or volume.loopback or volume._paths: print("\nPaths:") if volume.mountpoint: print(" - mountpoint {value}".format(value=volume.mountpoint)) if volume.loopback: print(" - loopback {value}".format(value=volume.loopback)) for k, v in volume._paths.items(): print(" - {name:<15} {value}".format(name=k, value=v)) if volume.info: print("\nAdditional information:") for k, v in volume.info.items(): print(" - {name:<15} {value}".format(name=k, value=v)) else: # displays a disk's details print(col("Details of disk {description}".format(description=disk.paths[0]), attrs=['bold'])) if disk.is_mounted: print(col("Currently mounted", 'green')) print("\nAttributes:") for attr in ['offset', 'block_size', 'read_write', 'disk_mounter']: print(" - {name:<15} {value}".format(name=attr, value=getattr(disk, attr))) if disk.mountpoint or disk.rwpath or disk._paths: print("\nPaths:") if disk.mountpoint: print(" - mountpoint {value}".format(value=disk.mountpoint)) if disk.rwpath: print(" - rwpath {value}".format(value=disk.rwpath)) for k, v in disk._paths.items(): print(" - {name:<15} {value}".format(name=k, value=v)) def _show_tree(self): col = get_coloring_func() for disk in self.parser.disks: print("- {index:<5} {type} {filename}".format( index=col("{:<5}".format(disk.index), 'green' if disk.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(disk.volumes.vstype), attrs=['dark']), filename=disk.paths[0] )) def _show_volume_system(volumes, level=0): level += 1 for i, v in enumerate(volumes): level_str = " "*level + ("└ " if i == len(volumes)-1 else "├ ") tp = v.volumes.vstype if v.fstype == 'volumesystem' else v.fstype if v.flag == 'alloc' else v.flag print("{level_str}{index} {type} {size:<10} {description}".format( level_str=level_str, index=col("{:<5}".format(v.index), 'green' if v.is_mounted else None, attrs=['bold']), type=col("{:<10}".format(tp), attrs=['dark']), description=v.get_description(with_index=False, with_size=False)[:30], size=v.get_formatted_size() )) _show_volume_system(v.volumes, level) _show_volume_system(disk.volumes) ###################################################################### # set command ###################################################################### def parser_set(self, parser): parser.description = "Modifies a property of a volume or disk. This is for advanced usage only." parser.add_argument('index', help='volume index', choices=self._get_all_indexes()) parser.add_argument('name', help='property name', choices=['size', 'offset', 'slot', 'flag', 'block_size', 'fstype', 'key', 'disk_mounter']) parser.add_argument('value', help='property value') def arg_set(self, args): col = get_coloring_func() volume, disk = self._get_by_index(args.index) if volume: if args.name in ['size', 'offset', 'block_size']: try: setattr(volume, args.name, int(args.value)) except ValueError: print(col("Invalid value provided for {}".format(args.name), 'red')) else: print(col("Updated value for {}".format(args.name), 'green')) elif args.name in ['slot', 'flag', 'fstype', 'key']: setattr(volume, args.name, args.value) print(col("Updated value for {}".format(args.name), 'green')) else: print(col("Property {} can't be set for a volume".format(args.name), 'red')) else: if args.name in ['offset', 'block_size']: try: setattr(disk, args.name, int(args.value)) except ValueError: print(col("Invalid value provided for {}".format(args.name), 'red')) else: print(col("Updated value for {}".format(args.name), 'green')) elif args.name in ['disk_mounter']: setattr(disk, args.name, args.value) print(col("Updated value for {}".format(args.name), 'green')) else: print(col("Property {} can't be set for a disk".format(args.name), 'red')) ###################################################################### # quit command ###################################################################### def do_save(self, arg): self.save() def do_EOF(self, arg): self.save() return True def do_quit(self, arg): """Quits the program.""" if self.saved: self.save() else: self.parser.clean() return True