pandaharvester/harvesterscripts/file_operation.py

0001 #!/usr/bin/env python
0002 #
0003 # Description:
0004 #   This script is meant to be able to run in any python environment since only builtin libraries are used
0005
0006
0007 import argparse
0008 import os
0009 import re
0010 import sys
0011 import tempfile
0012 import zlib
0013
0014 # === Command functions =========================================================
0015
0016
0017 def test(arguments):
0018     print("file_operation: test")
0019
0020
0021 # calculate adler32
0022
0023
0024 def adler32(arguments):
0025     file_name = arguments.file
0026     val = 1
0027     blockSize = 32 * 1024 * 1024
0028     with open(file_name, "rb") as fp:
0029         while True:
0030             data = fp.read(blockSize)
0031             if not data:
0032                 break
0033             val = zlib.adler32(data, val)
0034     if val < 0:
0035         val += 2**32
0036     retVal = hex(val)[2:10].zfill(8).lower()
0037     print(retVal)
0038
0039
0040 # write data into a temporary file; return the file name
0041
0042
0043 def write_tmpfile(arguments):
0044     tmpArgFile = tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=arguments.suffix, dir=arguments.dir)
0045     in_data = re.sub(r"\\\\", r"\\", arguments.data)
0046     in_data = re.sub(r"\\n", "\n", in_data)
0047     tmpArgFile.write(in_data)
0048     tmpArgFile.close()
0049     print(tmpArgFile.name)
0050
0051
0052 # remove file
0053
0054
0055 def remove_file(arguments):
0056     os.remove(arguments.path)
0057
0058
0059 # === Command map ===============================================================
0060
0061
0062 commandMap = {
0063     # test commands
0064     "test": test,
0065     # adler32 commands
0066     "adler32": adler32,
0067     "write_tmpfile": write_tmpfile,
0068     "remove_file": remove_file,
0069 }
0070
0071 # === Main ======================================================================
0072
0073
0074 def main():
0075     # main parser
0076     oparser = argparse.ArgumentParser(prog="file_operations.py", add_help=True)
0077     subparsers = oparser.add_subparsers()
0078
0079     # test command
0080     test_parser = subparsers.add_parser("test", help="for testing only")
0081     test_parser.set_defaults(which="test")
0082     # adler32 command
0083     adler32_parser = subparsers.add_parser("adler32", help="get adler32 checksum of the file")
0084     adler32_parser.set_defaults(which="adler32")
0085     adler32_parser.add_argument("file", type=str, action="store", metavar="<file>", help="file path")
0086     # write_tmpfile command
0087     write_tmpfile_parser = subparsers.add_parser("write_tmpfile", help="write data to a temporary file")
0088     write_tmpfile_parser.set_defaults(which="write_tmpfile")
0089     write_tmpfile_parser.add_argument("--suffix", type=str, action="store", metavar="<suffix>", default="xxx.tmp", help="name suffix of temporary file")
0090     write_tmpfile_parser.add_argument("--dir", type=str, action="store", metavar="<dir>", default="/tmp", help="directory of temorary file")
0091     write_tmpfile_parser.add_argument("data", type=str, action="store", metavar="<data>", help="data to write in temporary file")
0092     # remove_file command
0093     remove_file_parser = subparsers.add_parser("remove_file", help="remove a file")
0094     remove_file_parser.set_defaults(which="remove_file")
0095     remove_file_parser.add_argument("path", type=str, action="store", metavar="<path>", help="file path")
0096
0097     # start parsing
0098     if len(sys.argv) == 1:
0099         oparser.print_help()
0100         sys.exit(1)
0101     arguments = oparser.parse_args(sys.argv[1:])
0102     # Run command functions
0103     try:
0104         command = commandMap.get(arguments.which)
0105     except AttributeError:
0106         oparser.print_help()
0107         sys.exit(1)
0108     try:
0109         result = command(arguments)
0110         sys.exit(result)
0111     except Exception as e:
0112         sys.stderr.write(f"{e.__class__.__name__}: {e}")
0113         sys.exit(1)
0114
0115
0116 if __name__ == "__main__":
0117     main()