tfscript/src/tfscript/cli.py

"""
Command line module for making Team Fortress 2 macro scripts from
YAML source code.
"""

__all__     = ['parseFile']
__author__  = "Nicholas Hope <tfscript@nickhope.world"
__date__    = "26 August 2022"
__version__ = "1.0"
__copyright__ = "Copyright © 2022 Nicholas Hope. See LICENSE for details."

# Standard libraries
import sys
import os
import argparse
import tempfile
import yaml
from platform import system as getOS, release

try:
    import winreg
except ModuleNotFoundError:
    # Not running on windows
    pass

# Local libraries
import tfscript
from tfscript import verify

args = {}
targetDir = ""

def parseFile(inputFile):
    """Parse, verify, and do the conversion."""
    config = yaml.safe_load(inputFile)

    # See verify.py
    config = tfscript.verify.verifyConfig(config)
    if "errors" in config:
        for e in config["errors"]:
            print(e,file=sys.stderr)
    else:
        parseConfig(config)

def writeOutput(scriptString, className):
    """Given the string of stuff to write, write it out to the given handle."""
    global args
    global targetDir
    chunksize = 2**20 # 1Mb maximum cfg file size
    chunk = 1
    # If the string is more than 1048576 bytes, we need divide it into files that each
    # are less than 1048576 bytes
    chunksneeded = int( 1 + len(scriptString) / chunksize )
    if args.debug:
        print( f'DEBUG: need {chunksneeded} files for {className}')

    if( chunksneeded == 1):
        # If it can be done in one chunk, do it in one chunk.
        outfile = tempfile.NamedTemporaryFile( prefix=className, delete=False )
        outfile.write(scriptString.encode("utf8"))
        outfile.close()
        if( args.dry_run != True ):
            os.replace(outfile.name, f'{targetDir}/{className}_script_{chunk:02d}.cfg')
            if args.debug:
                print( f'DEBUG: Created {targetDir}/{className}_script_{chunk:02d}.cfg')
        else:
            if args.debug:
                print( f'DEBUG: {outfile.name} would be {targetDir}/{className}_script_{chunk:02d}.cfg')
    else:
        # Gotta do it in multiple chunks
        classLines = scriptString.splitlines()
        execString = f'exec {className}_script_{chunk:02d}'.encode("utf8")
        # extra 4 bytes is just a little buffer so we don't get exactly chunksize bytes
        reservedSpace = len(execString) + 4
        n = 0
        pieces = {}
        while( chunk <= chunksneeded ):
            outfile = tempfile.NamedTemporaryFile( prefix=className, delete=False )
            pieces[outfile.name] = f'{targetDir}/{className}_script_{chunk:02d}.cfg'
            byteswritten = 0
            while( n < len(classLines) and (byteswritten + len(classLines[n]) + reservedSpace) < chunksize ):
                line = classLines[n].encode("utf8") + os.linesep.encode("utf8")
                outfile.write(line)
                byteswritten += len(line)
                n+=1
            if( chunk < chunksneeded ):
                line = f'exec {className}_script_{chunk+1:02d}'.encode("utf8") + os.linesep.encode("utf8")
                outfile.write(line)
                byteswritten += len(line)
            outfile.close()
            
            if args.debug:
                print( f'DEBUG: Wrote {byteswritten} bytes to {className} ({chunk}/{chunksneeded})' )
            chunk += 1
        for tmpname, realname in pieces.items():
            if( args.dry_run ):
                if( args.debug ):
                    print( f'DEBUG: {outfile.name} would be {targetDir}/{className}_script_{chunk:02d}.cfg')
            else:
                os.replace(tmpname, realname)
                print( f'DEBUG: Created {targetDir}/{className}_script_{chunk:02d}.cfg')

def parseConfig(config):
    """With validated data structure, write out all the files."""
    global targetDir
    # Make sure the target exists before we try to use it
    if os.path.isdir( targetDir ) == False:
        try:
            os.mkdir( targetDir )
            if args.debug:
                print( f'DEBUG: created {targetDir}')
        except Exception as fileExcept:
            print( f'WARN: Failed to create {targetDir}: {fileExcept.strerror}\nUsing current directory instead.' )
            targetDir = '.'
    for currentClass in config:
        classDict = config[currentClass]
        stringToWrite = tfscript.makeCFG(classDict)
        writeOutput(stringToWrite, currentClass)

def parseCLI():
    # Handle command line
    parser = argparse.ArgumentParser(
        description="Parse YAML file and produce TF2 config script."
    )
    parser.add_argument( '-d', '--debug', action='store_true',
        help="Enable debugging messages.")
    parser.add_argument( '-n', '--dry-run', action='store_true',
        help="Parse input file, but don't write anything.")
    parser.add_argument( '-f', '--force', action='store_true',
        help="Force tfscript to continue until catastrophic failure")
    parser.add_argument( '-D', '--directory', action='store', type=str,
        help="Change output directory")
    # positional argument: first non-hyphenated argument is input file
    parser.add_argument( 'infile', type=argparse.FileType('r'),
        help='File containing YAML to convert.')
    return parser

def main():
    """
    Command line interface.
    """
    global args
    global targetDir
    parser = parseCLI()

    args = parser.parse_args()
    if args.directory is not None:
        targetDir = args.directory
    else:
        systemName = getOS()
        if systemName == "Darwin":
            if float( '.'.join(release().split('.')[0:2]) ) >= 10.15:
                if not args.force:
                    print(
                        "As of macOS Catalina (v10.15), 32-bit applications "
                        "such as tf2 do not work, so tfscript does not function",
                        file=sys.stderr
                    )
            else:
                targetDir = os.path.expanduser("~/Library/Application Support/Steam/")

        elif systemName == "Windows":
            # oh god why do we have to use the registry
            accessReg = winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE)
            accessKey = winreg.OpenKey(accessReg, "SOFTWARE\\WOW6432Node\\Valve\\Steam\\")
            targetDir = winreg.QueryValue(accessKey, "InstallPath")

        elif systemName == "Linux":
            targetDir = os.path.expanduser("~/.local/Steam")

        elif systemName == "Java":
            print("Java-based OSes are not supported yet by tfscript.", file=sys.stderr)

        if targetDir != "":
            # Supported OS: add steamapps path
            if targetDir[-1] != '/':
                targetDir += '/'
            targetDir += "steamapps/common/Team Fortress 2/tf/cfg"
        elif args.force:
            # Unsupported OS but -f specified
            targetDir = os.path.expanduser('.')
        else:
            # Unsupported OS and not forced to continue
            return 2

    parseFile(args.infile)
    return 0

if __name__ == "__main__":
    exit(main())