BaseTools: Convert Split tool to python

[mirror_edk2.git] / BaseTools / Source / Python / Split / Split.py

# @file\r
#  Split a file into two pieces at the request offset.\r
#\r
#  Copyright (c) 2021, Intel Corporation. All rights reserved.<BR>\r
#\r
#  SPDX-License-Identifier: BSD-2-Clause-Patent\r
#\r
##\r
\r
# Import Modules\r
#\r
import argparse\r
import os\r
import io\r
import shutil\r
import logging\r
import sys\r
import tempfile\r
\r
parser = argparse.ArgumentParser(description='''\r
SplitFile creates two Binary files either in the same directory as the current working directory or in the specified directory.\r
''')\r
parser.add_argument("-f", "--filename", dest="inputfile",\r
                    required=True, help="The input file to split tool.")\r
parser.add_argument("-s", "--split", dest="position",\r
                    required=True, help="The number of bytes in the first file. The valid format are HEX, Decimal and Decimal[KMG].")\r
parser.add_argument("-p", "--prefix",  dest="output",\r
                    help="The output folder.")\r
parser.add_argument("-o", "--firstfile",  help="The first file name")\r
parser.add_argument("-t", "--secondfile",  help="The second file name")\r
parser.add_argument("--version", action="version", version='%(prog)s Version 2.0',\r
                    help="Print debug information.")\r
\r
group = parser.add_mutually_exclusive_group()\r
group.add_argument("-v", "--verbose", action="store_true",\r
                   help="Print debug information.")\r
group.add_argument("-q", "--quiet", action="store_true",\r
                   help="Disable all messages except fatal errors")\r
\r
SizeDict = {\r
    "K": 1024,\r
    "M": 1024*1024,\r
    "G": 1024*1024*1024\r
}\r
\r
\r
def GetPositionValue(position):\r
    '''\r
    Parse the string of the argument position and return a decimal number.\r
    The valid position formats are\r
    1. HEX\r
    e.g. 0x1000 or 0X1000\r
    2. Decimal\r
    e.g. 100\r
    3. Decimal[KMG]\r
    e.g. 100K or 100M or 100G or 100k or 100m or 100g\r
    '''\r
    logger = logging.getLogger('Split')\r
    PosVal = 0\r
    header = position[:2].upper()\r
    tailer = position[-1].upper()\r
\r
    try:\r
        if tailer in SizeDict:\r
            PosVal = int(position[:-1]) * SizeDict[tailer]\r
        else:\r
            if header == "0X":\r
                PosVal = int(position, 16)\r
            else:\r
                PosVal = int(position)\r
    except Exception as e:\r
        logger.error(\r
            "The parameter %s format is incorrect. The valid format is HEX, Decimal and Decimal[KMG]." % position)\r
        raise(e)\r
\r
    return PosVal\r
\r
\r
def getFileSize(filename):\r
    '''\r
    Read the input file and return the file size.\r
    '''\r
    logger = logging.getLogger('Split')\r
    length = 0\r
    try:\r
        with open(filename, "rb") as fin:\r
            fin.seek(0, io.SEEK_END)\r
            length = fin.tell()\r
    except Exception as e:\r
        logger.error("Access file failed: %s", filename)\r
        raise(e)\r
\r
    return length\r
\r
\r
def splitFile(inputfile, position, outputdir=None, outputfile1=None, outputfile2=None):\r
    '''\r
    Split the inputfile into outputfile1 and outputfile2 from the position.\r
    '''\r
    logger = logging.getLogger('Split')\r
\r
    inputfile = os.path.abspath(inputfile)\r
    workspace = os.path.dirname(inputfile)\r
    if not os.path.exists(inputfile):\r
        logger.error("File Not Found: %s" % inputfile)\r
        raise(Exception)\r
\r
    if outputfile1 and outputfile2 and outputfile1 == outputfile2:\r
        logger.error(\r
            "The firstfile and the secondfile can't be the same: %s" % outputfile1)\r
        raise(Exception)\r
\r
    if not outputdir:\r
        outputdir = workspace\r
    elif not os.path.isabs(outputdir):\r
        outputdir = os.path.join(workspace, outputdir)\r
\r
    # Create dir for the output files\r
    try:\r
        if not outputfile1:\r
            outputfile1 = os.path.abspath(os.path.join(\r
                outputdir, "{}1".format(os.path.basename(inputfile))))\r
        else:\r
            outputfile1 = os.path.abspath(os.path.join(outputdir, outputfile1))\r
        outputdir = os.path.dirname(outputfile1)\r
        if not os.path.exists(outputdir):\r
            os.makedirs(outputdir)\r
\r
        if not outputfile2:\r
            outputfile2 = os.path.abspath(os.path.join(\r
                outputdir, "{}2".format(os.path.basename(inputfile))))\r
        else:\r
            outputfile2 = os.path.abspath(os.path.join(outputdir, outputfile2))\r
        outputdir = os.path.dirname(outputfile2)\r
        if not os.path.exists(outputdir):\r
            os.makedirs(outputdir)\r
    except Exception as e:\r
        logger.error("Can't make dir: %s" % outputdir)\r
        raise(e)\r
\r
    if position <= 0:\r
        if outputfile2 != inputfile:\r
            shutil.copy2(inputfile, outputfile2)\r
        with open(outputfile1, "wb") as fout:\r
            fout.write(b'')\r
    else:\r
        inputfilesize = getFileSize(inputfile)\r
        if position >= inputfilesize:\r
            if outputfile1 != inputfile:\r
                shutil.copy2(inputfile, outputfile1)\r
            with open(outputfile2, "wb") as fout:\r
                fout.write(b'')\r
        else:\r
            try:\r
                tempdir = tempfile.mkdtemp()\r
                tempfile1 = os.path.join(tempdir, "file1.bin")\r
                tempfile2 = os.path.join(tempdir, "file2.bin")\r
                with open(inputfile, "rb") as fin:\r
                    content1 = fin.read(position)\r
                    with open(tempfile1, "wb") as fout1:\r
                        fout1.write(content1)\r
\r
                    content2 = fin.read(inputfilesize - position)\r
                    with open(tempfile2, "wb") as fout2:\r
                        fout2.write(content2)\r
                shutil.copy2(tempfile1, outputfile1)\r
                shutil.copy2(tempfile2, outputfile2)\r
            except Exception as e:\r
                logger.error("Split file failed")\r
                raise(e)\r
            finally:\r
                if os.path.exists(tempdir):\r
                    shutil.rmtree(tempdir)\r
\r
\r
def main():\r
    args = parser.parse_args()\r
    status = 0\r
\r
    logger = logging.getLogger('Split')\r
    if args.quiet:\r
        logger.setLevel(logging.CRITICAL)\r
    if args.verbose:\r
        logger.setLevel(logging.DEBUG)\r
\r
    lh = logging.StreamHandler(sys.stdout)\r
    lf = logging.Formatter("%(levelname)-8s: %(message)s")\r
    lh.setFormatter(lf)\r
    logger.addHandler(lh)\r
\r
    try:\r
        position = GetPositionValue(args.position)\r
        splitFile(args.inputfile, position, args.output,\r
                  args.firstfile, args.secondfile)\r
    except Exception as e:\r
        status = 1\r
\r
    return status\r
\r
\r
if __name__ == "__main__":\r
    exit(main())\r