<commit_before>#!/usr/bin/env python
import argparse
import glob
import sys
import os

from imagemounter import util, ImageParser, __version__
from termcolor import colored


def main():
    class MyParser(argparse.ArgumentParser):
        def error(self, message):
            sys.stderr.write('error: {0}\n'.format(message))
            self.print_help()
            sys.exit(2)

    class CleanAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            commands = ImageParser.force_clean(False)
            if not commands:
                print "[+] Nothing to do"
                parser.exit()
            print "[!] --clean will rigorously clean anything that looks like a mount or volume group originating " \
                  "from this utility. You may regret using this if you have other mounts or volume groups that are " \
                  "similarly named. The following commands will be executed:"
            for c in commands:
                print "    {0}".format(c)
            try:
                raw_input(">>> Press [enter] to continue or ^C to cancel... ")
                ImageParser.force_clean()
            except KeyboardInterrupt:
                print "\n[-] Aborted."
            parser.exit()

    parser = MyParser(description=u'Utility to mount volumes in Encase and dd images locally.')
    parser.add_argument('images', nargs='+',
                        help='path(s) to the image(s) that you want to mount; generally just the first file (e.g. '
                             'the .E01 or .001 file) or the folder containing the files is enough in the cast of '
                             'split files.')
    parser.add_argument('--version', action='version', version=__version__, help='display version and exit')
    parser.add_argument('--clean', action=CleanAction, nargs=0,
                        help='try to rigorously clean anything that resembles traces from previous runs of '
                             'this utility')
    parser.add_argument('-c', '--color', action='store_true', default=False, help='colorize the output')
    parser.add_argument('-w', '--wait', action='store_true', default=False, help='pause on some additional warnings')
    parser.add_argument('-r', '--reconstruct', action='store_true', default=False,
                        help='attempt to reconstruct the full filesystem tree; implies -s and mounts all partitions '
                             'at once')
    parser.add_argument('-rw', '--read-write', action='store_true', default=False,
                        help='mount image read-write by creating a local write-cache file in a temp directory; '
                             'implies --method=xmount')
    parser.add_argument('-s', '--stats', action='store_true', default=True,
                        help='show limited information from fsstat, which will slow down mounting and may cause '
                             'random issues such as partitions being unreadable (default)')
    parser.add_argument('--no-stats', action='store_true', default=False,
                        help='do not show limited information from fsstat')
    parser.add_argument('-m', '--method', choices=['xmount', 'affuse', 'ewfmount', 'auto'], default='auto',
                        help='use other tool to mount the initial images; results may vary between methods and if '
                             'something doesn\'t work, try another method (default: auto)')
    parser.add_argument('-md', '--mountdir', default=None,
                        help='specify other directory for volume mountpoints')
    parser.add_argument('-p', '--pretty', action='store_true', default=False,
                        help='use pretty names for mount points; useful in combination with --mountdir and does not '
                             'provide a fallback when the pretty mount name is unavailable')
    parser.add_argument('-vs', '--vstype', choices=['detect', 'dos', 'bsd', 'sun', 'mac', 'gpt', 'dbfiller', 'any'],
                        default="detect", help='specify type of volume system (partition table); if you don\'t know, '
                                               'use "detect" to try to detect, or "any" to loop over all VS types and '
                                               'use whatever works, which may produce unexpected results (default: '
                                               'detect)')
    parser.add_argument('-fs', '--fstype', choices=['ext', 'ufs', 'ntfs', 'lvm', 'unknown'], default=None,
                        help="specify fallback type of the filesystem, which is used when it could not be detected or "
                             "is unsupported; use unknown to mount without specifying type")
    parser.add_argument('-fsf', '--fsforce', action='store_true', default=False,
                        help="force the use of the filesystem type specified with --fstype for all volumes")
    parser.add_argument('-k', '--keep', action='store_true', default=False,
                        help='keep volumes mounted after program exits')
    parser.add_argument('-v', '--verbose', action='store_true', default=False, help='enable verbose output')
    args = parser.parse_args()

    # Check some prerequisites
    if os.geteuid():  # Not run as root
        print u'[-] This program needs to be ran as root! Requesting elevation... '
        os.execvp('sudo', ['sudo'] + sys.argv)
        #sys.exit(1)

    if not args.color:
        #noinspection PyUnusedLocal,PyShadowingNames
        def col(s, *args, **kwargs):
            return s
    else:
        col = colored

    if args.no_stats:
        args.stats = False

    if args.method not in ('xmount', 'auto') and args.read_write:
        print "[-] {0} does not support mounting read-write! Will mount read-only.".format(args.method)
        args.read_write = False

    if args.method != 'auto' and not util.command_exists(args.method):
        print "[-] {0} is not installed!".format(args.method)
        sys.exit(1)

    if args.method == 'auto' and not any(map(util.command_exists, ('xmount', 'affuse', 'ewfmount'))):
        print "[-] No tools installed to mount the image base!"
        sys.exit(1)

    if args.reconstruct:  # Reconstruct implies use of fsstat
        args.stats = True

    if args.stats and not util.command_exists('fsstat'):
        print "[-] To obtain stats, the fsstat command is used (part of sleuthkit package), but is not installed. " \
              "Stats will not be obtained during this session."
        args.stats = False

        if args.reconstruct:
            print "[-] Reconstruction is now impossible!"
            sys.exit(1)

    if args.fstype and not args.fsforce:
        print "[!] You are using the file system type {0} as fallback. This may cause unexpected results."\
            .format(args.fstype)
    elif args.fstype and args.fsforce:
        print "[!] You are forcing the file system type to {0}. This may cause unexpected results.".format(args.fstype)
    elif not args.fstype and args.fsforce:
        print "[-] You are forcing a file system type, but have not specified the type to use. Ignoring force."
        args.fsforce = False

    if args.vstype == 'any':
        print "[!] You are using the 'any' volume system type. This may cause unexpected results. It is recommended " \
              "to use either 'detect' or explicitly specify the correct volume system. However, 'any' may provide " \
              "some hints on the volume system to use (e.g. GPT mounted as DOS lists a GPT safety partition)."

    # Enumerate over all images in the CLI
    for num, image in enumerate(args.images):
        # If is a directory, find a E01 file in the directory
        if os.path.isdir(image):
            for f in glob.glob(os.path.join(image, '*.[E0]01')):
                image = f
                break
            else:
                print col("[-] {0} is a directory, aborting!".format(image), "red")
                break

        elif not os.path.exists(image):
            print col("[-] Image {0} does not exist, aborting!".format(image), "red")
            break

        try:
            p = ImageParser(image, **vars(args))
            print u'[+] Mounting image {0} using {1}...'.format(p.paths[0], p.method)

            # Mount the base image using the preferred method
            if not p.mount_base():
                print col("[-] Failed mounting base image.", "red")
                continue

            if args.read_write:
                print u'[+] Created read-write cache at {0}'.format(p.rwpath)
            print u'[+] Mounted raw image [{num}/{total}], now mounting volumes...'.format(num=num + 1,
                                                                                           total=len(args.images))

            sys.stdout.write("[+] Mounting partition 0...\r")
            sys.stdout.flush()<commit_msg>Small changes in argument parsing.<commit_after>#!/usr/bin/env python
import argparse
import glob
import sys
import os

from imagemounter import util, ImageParser, __version__
from termcolor import colored


def main():
    class MyParser(argparse.ArgumentParser):
        def error(self, message):
            sys.stderr.write('error: {0}\n'.format(message))
            self.print_help()
            sys.exit(2)

    class CleanAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            commands = ImageParser.force_clean(False)
            if not commands:
                print "[+] Nothing to do"
                parser.exit()
            print "[!] --clean will rigorously clean anything that looks like a mount or volume group originating " \
                  "from this utility. You may regret using this if you have other mounts or volume groups that are " \
                  "similarly named. The following commands will be executed:"
            for c in commands:
                print "    {0}".format(c)
            try:
                raw_input(">>> Press [enter] to continue or ^C to cancel... ")
                ImageParser.force_clean()
            except KeyboardInterrupt:
                print "\n[-] Aborted."
            parser.exit()

    parser = MyParser(description=u'Utility to mount volumes in Encase and dd images locally.')
    parser.add_argument('images', nargs='+',
                        help='path(s) to the image(s) that you want to mount; generally just the first file (e.g. '
                             'the .E01 or .001 file) or the folder containing the files is enough in the case of '
                             'split files.')
    parser.add_argument('--version', action='version', version=__version__, help='display version and exit')
    parser.add_argument('--clean', action=CleanAction, nargs=0,
                        help='try to rigorously clean anything that resembles traces from previous runs of '
                             'this utility')
    parser.add_argument('-c', '--color', action='store_true', default=False, help='colorize the output')
    parser.add_argument('-w', '--wait', action='store_true', default=False, help='pause on some additional warnings')
    parser.add_argument('-k', '--keep', action='store_true', default=False,
                        help='keep volumes mounted after program exits')
    parser.add_argument('-r', '--reconstruct', action='store_true', default=False,
                        help='attempt to reconstruct the full filesystem tree; implies -s and mounts all partitions '
                             'at once')
    parser.add_argument('-rw', '--read-write', action='store_true', default=False,
                        help='mount image read-write by creating a local write-cache file in a temp directory; '
                             'implies --method=xmount')
    parser.add_argument('-s', '--stats', action='store_true', default=False,
                        help='show limited information from fsstat, which will slow down mounting and may cause '
                             'random issues such as partitions being unreadable (default)')
    parser.add_argument('-n', '--no-stats', action='store_true', default=False,
                        help='do not show limited information from fsstat')
    parser.add_argument('-m', '--method', choices=['xmount', 'affuse', 'ewfmount', 'auto'], default='auto',
                        help='use other tool to mount the initial images; results may vary between methods and if '
                             'something doesn\'t work, try another method (default: auto)')
    parser.add_argument('-md', '--mountdir', default=None,
                        help='specify other directory for volume mountpoints')
    parser.add_argument('-p', '--pretty', action='store_true', default=False,
                        help='use pretty names for mount points; useful in combination with --mountdir and does not '
                             'provide a fallback when the pretty mount name is unavailable')
    parser.add_argument('--vstype', choices=['detect', 'dos', 'bsd', 'sun', 'mac', 'gpt', 'dbfiller', 'any'],
                        default="detect", help='specify type of volume system (partition table); if you don\'t know, '
                                               'use "detect" to try to detect, or "any" to loop over all VS types and '
                                               'use whatever works, which may produce unexpected results (default: '
                                               'detect)')
    parser.add_argument('--fstype', choices=['ext', 'ufs', 'ntfs', 'lvm', 'unknown'], default=None,
                        help="specify fallback type of the filesystem, which is used when it could not be detected or "
                             "is unsupported; use unknown to mount without specifying type")
    parser.add_argument('--fsforce', action='store_true', default=False,
                        help="force the use of the filesystem type specified with --fstype for all volumes")
    parser.add_argument('-v', '--verbose', action='store_true', default=False, help='enable verbose output')
    args = parser.parse_args()

    # Check some prerequisites
    if os.geteuid():  # Not run as root
        print u'[-] This program needs to be ran as root! Requesting elevation... '
        os.execvp('sudo', ['sudo'] + sys.argv)
        #sys.exit(1)

    if not args.color:
        #noinspection PyUnusedLocal,PyShadowingNames
        def col(s, *args, **kwargs):
            return s
    else:
        col = colored

    # Always assume stats, except when --no-stats is present, and --stats is not.
    if not args.stats and args.no_stats:
        args.stats = False
    else:
        args.stats = True

    if args.method not in ('xmount', 'auto') and args.read_write:
        print "[-] {0} does not support mounting read-write! Will mount read-only.".format(args.method)
        args.read_write = False

    if args.method != 'auto' and not util.command_exists(args.method):
        print "[-] {0} is not installed!".format(args.method)
        sys.exit(1)
    elif args.method == 'auto' and not any(map(util.command_exists, ('xmount', 'affuse', 'ewfmount'))):
        print "[-] No tools installed to mount the image base!"
        sys.exit(1)

    if args.reconstruct and not args.stats:  # Reconstruct implies use of fsstat
        print "[-] You explicitly disabled stats, but --reconstruct implies the use of stats. Stats are re-enabled."
        args.stats = True

    if args.stats and not util.command_exists('fsstat'):
        print "[-] The fsstat command (part of sleuthkit package) is required to obtain stats, but is not installed. " \
              "Stats can not be obtained during this session."
        args.stats = False

        if args.reconstruct:
            print "[-] Reconstruction requires stats to be obtained, but stats can not be enabled."
            sys.exit(1)

    if args.fstype and not args.fsforce:
        print "[!] You are using the file system type {0} as fallback. This may cause unexpected results."\
            .format(args.fstype)
    elif args.fstype and args.fsforce:
        print "[!] You are forcing the file system type to {0}. This may cause unexpected results.".format(args.fstype)
    elif not args.fstype and args.fsforce:
        print "[-] You are forcing a file system type, but have not specified the type to use. Ignoring force."
        args.fsforce = False

    if args.vstype == 'any':
        print "[!] You are using the 'any' volume system type. This may cause unexpected results. It is recommended " \
              "to use either 'detect' or explicitly specify the correct volume system. However, 'any' may provide " \
              "some hints on the volume system to use (e.g. GPT mounted as DOS lists a GPT safety partition)."

    # Enumerate over all images in the CLI
    for num, image in enumerate(args.images):
        # If is a directory, find a E01 file in the directory
        if os.path.isdir(image):
            for f in glob.glob(os.path.join(image, '*.[E0]01')):
                image = f
                break
            else:
                print col("[-] {0} is a directory not containing a .001 or .E01 file, aborting!".format(image), "red")
                break

        elif not os.path.exists(image):
            print col("[-] Image {0} does not exist, aborting!".format(image), "red")
            break

        try:
            p = ImageParser(image, **vars(args))
            print u'[+] Mounting image {0} using {1}...'.format(p.paths[0], p.method)

            # Mount the base image using the preferred method
            if not p.mount_base():
                print col("[-] Failed mounting base image.", "red")
                continue

            if args.read_write:
                print u'[+] Created read-write cache at {0}'.format(p.rwpath)
            print u'[+] Mounted raw image [{num}/{total}], now mounting volumes...'.format(num=num + 1,
                                                                                           total=len(args.images))

            sys.stdout.write("[+] Mounting partition 0...\r")
            sys.stdout.flush()
{'fraction_non_alphanumeric': 0.11176245000821783, 'fraction_numerical': 0.0035610584561441955, 'mean_word_length': 2.0850092952509716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1049198', 'n_tokens_mistral': 4839, 'n_tokens_neox': 4644, 'n_words': 1712}