BaseTools/Source/Python/Split/Split.py

   1 # @file
   2 #  Split a file into two pieces at the request offset.
   3 #
   4 #  Copyright (c) 2021, Intel Corporation. All rights reserved.<BR>
   5 #
   6 #  SPDX-License-Identifier: BSD-2-Clause-Patent
   7 #
   8 ##
   9
  10 # Import Modules
  11 #
  12 import argparse
  13 import os
  14 import io
  15 import shutil
  16 import logging
  17 import sys
  18 import tempfile
  19
  20 parser = argparse.ArgumentParser(description='''
  21 SplitFile creates two Binary files either in the same directory as the current working directory or in the specified directory.
  22 ''')
  23 parser.add_argument("-f", "--filename", dest="inputfile",
  24                     required=True, help="The input file to split tool.")
  25 parser.add_argument("-s", "--split", dest="position",
  26                     required=True, help="The number of bytes in the first file. The valid format are HEX, Decimal and Decimal[KMG].")
  27 parser.add_argument("-p", "--prefix",  dest="output",
  28                     help="The output folder.")
  29 parser.add_argument("-o", "--firstfile",  help="The first file name")
  30 parser.add_argument("-t", "--secondfile",  help="The second file name")
  31 parser.add_argument("--version", action="version", version='%(prog)s Version 2.0',
  32                     help="Print debug information.")
  33
  34 group = parser.add_mutually_exclusive_group()
  35 group.add_argument("-v", "--verbose", action="store_true",
  36                    help="Print debug information.")
  37 group.add_argument("-q", "--quiet", action="store_true",
  38                    help="Disable all messages except fatal errors")
  39
  40 SizeDict = {
  41     "K": 1024,
  42     "M": 1024*1024,
  43     "G": 1024*1024*1024
  44 }
  45
  46
  47 def GetPositionValue(position):
  48     '''
  49     Parse the string of the argument position and return a decimal number.
  50     The valid position formats are
  51     1. HEX
  52     e.g. 0x1000 or 0X1000
  53     2. Decimal
  54     e.g. 100
  55     3. Decimal[KMG]
  56     e.g. 100K or 100M or 100G or 100k or 100m or 100g
  57     '''
  58     logger = logging.getLogger('Split')
  59     PosVal = 0
  60     header = position[:2].upper()
  61     tailer = position[-1].upper()
  62
  63     try:
  64         if tailer in SizeDict:
  65             PosVal = int(position[:-1]) * SizeDict[tailer]
  66         else:
  67             if header == "0X":
  68                 PosVal = int(position, 16)
  69             else:
  70                 PosVal = int(position)
  71     except Exception as e:
  72         logger.error(
  73             "The parameter %s format is incorrect. The valid format is HEX, Decimal and Decimal[KMG]." % position)
  74         raise(e)
  75
  76     return PosVal
  77
  78
  79 def getFileSize(filename):
  80     '''
  81     Read the input file and return the file size.
  82     '''
  83     logger = logging.getLogger('Split')
  84     length = 0
  85     try:
  86         with open(filename, "rb") as fin:
  87             fin.seek(0, io.SEEK_END)
  88             length = fin.tell()
  89     except Exception as e:
  90         logger.error("Access file failed: %s", filename)
  91         raise(e)
  92
  93     return length
  94
  95 def getoutputfileabs(inputfile, prefix, outputfile,index):
  96     inputfile = os.path.abspath(inputfile)
  97     if outputfile is None:
  98         if prefix is None:
  99             outputfileabs = os.path.join(os.path.dirname(inputfile), "{}{}".format(os.path.basename(inputfile),index))
 100         else:
 101             if os.path.isabs(prefix):
 102                 outputfileabs = os.path.join(prefix, "{}{}".format(os.path.basename(inputfile),index))
 103             else:
 104                 outputfileabs = os.path.join(os.getcwd(), prefix, "{}{}".format(os.path.basename(inputfile),index))
 105     elif not os.path.isabs(outputfile):
 106         if prefix is None:
 107             outputfileabs = os.path.join(os.getcwd(), outputfile)
 108         else:
 109             if os.path.isabs(prefix):
 110                 outputfileabs = os.path.join(prefix, outputfile)
 111             else:
 112                 outputfileabs = os.path.join(os.getcwd(), prefix, outputfile)
 113     else:
 114         outputfileabs = outputfile
 115     return outputfileabs
 116
 117 def splitFile(inputfile, position, outputdir=None, outputfile1=None, outputfile2=None):
 118     '''
 119     Split the inputfile into outputfile1 and outputfile2 from the position.
 120     '''
 121     logger = logging.getLogger('Split')
 122
 123     if not os.path.exists(inputfile):
 124         logger.error("File Not Found: %s" % inputfile)
 125         raise(Exception)
 126
 127     if outputfile1 and outputfile2 and outputfile1 == outputfile2:
 128         logger.error(
 129             "The firstfile and the secondfile can't be the same: %s" % outputfile1)
 130         raise(Exception)
 131
 132     # Create dir for the output files
 133     try:
 134
 135         outputfile1 = getoutputfileabs(inputfile, outputdir, outputfile1,1)
 136         outputfolder = os.path.dirname(outputfile1)
 137         if not os.path.exists(outputfolder):
 138             os.makedirs(outputfolder)
 139
 140         outputfile2 = getoutputfileabs(inputfile, outputdir, outputfile2,2)
 141         outputfolder = os.path.dirname(outputfile2)
 142         if not os.path.exists(outputfolder):
 143             os.makedirs(outputfolder)
 144
 145     except Exception as e:
 146         logger.error("Can't make dir: %s" % outputfolder)
 147         raise(e)
 148
 149     if position <= 0:
 150         if outputfile2 != os.path.abspath(inputfile):
 151             shutil.copyfile(os.path.abspath(inputfile), outputfile2)
 152         with open(outputfile1, "wb") as fout:
 153             fout.write(b'')
 154     else:
 155         inputfilesize = getFileSize(inputfile)
 156         if position >= inputfilesize:
 157             if outputfile1 != os.path.abspath(inputfile):
 158                 shutil.copyfile(os.path.abspath(inputfile), outputfile1)
 159             with open(outputfile2, "wb") as fout:
 160                 fout.write(b'')
 161         else:
 162             try:
 163                 tempdir = tempfile.mkdtemp()
 164                 tempfile1 = os.path.join(tempdir, "file1.bin")
 165                 tempfile2 = os.path.join(tempdir, "file2.bin")
 166                 with open(inputfile, "rb") as fin:
 167                     content1 = fin.read(position)
 168                     with open(tempfile1, "wb") as fout1:
 169                         fout1.write(content1)
 170
 171                     content2 = fin.read(inputfilesize - position)
 172                     with open(tempfile2, "wb") as fout2:
 173                         fout2.write(content2)
 174                 shutil.copyfile(tempfile1, outputfile1)
 175                 shutil.copyfile(tempfile2, outputfile2)
 176             except Exception as e:
 177                 logger.error("Split file failed")
 178                 raise(e)
 179             finally:
 180                 if os.path.exists(tempdir):
 181                     shutil.rmtree(tempdir)
 182
 183
 184 def main():
 185     args = parser.parse_args()
 186     status = 0
 187
 188     logger = logging.getLogger('Split')
 189     if args.quiet:
 190         logger.setLevel(logging.CRITICAL)
 191     if args.verbose:
 192         logger.setLevel(logging.DEBUG)
 193
 194     lh = logging.StreamHandler(sys.stdout)
 195     lf = logging.Formatter("%(levelname)-8s: %(message)s")
 196     lh.setFormatter(lf)
 197     logger.addHandler(lh)
 198
 199     try:
 200         position = GetPositionValue(args.position)
 201         splitFile(args.inputfile, position, args.output,
 202                   args.firstfile, args.secondfile)
 203     except Exception as e:
 204         status = 1
 205
 206     return status
 207
 208
 209 if __name__ == "__main__":
 210     exit(main())