extract-bc

#!/usr/bin/env python
"""
This tool can be used two ways.

If the passed in file is a binary executable it
will extract the .llvm_bc section from the provided ELF object and reassemble
it into an actual bitcode file.  The ELF section contains absolute paths to all
of its constituent bitcode files.  This utility reads the section and links
together all of the named bitcode files.

If the passed in file is a static library it will extract the constituent
ELF objects and read their .llvm_bc sections and create a LLVM Bitcode
archive from the bitcode files.

"""

import os
import sys
from subprocess import *
from driver.utils import llvmCompilerPathEnv
from driver.popenwrapper import Popen
from driver.utils import elfSectionName
from driver.utils import FileType

import logging
import pprint
import driver.logconfig
import tempfile
import shutil

# Python 2 does not have exceptions automatically
# imported whereas python 3 does. Handle this
try:
    dir(UnicodeDecodeError)
except NameError:
  import exceptions

bitCodeArchiveExtension='bca'
moduleExtension='bc'

# Use objdump on the provided binary; parse out the fields
# to find the given section.  Return the size and offset of
# that section (in bytes)
def getSectionSizeAndOffset(sectionName, filename):
    objdumpCmd = ['objdump', '-h', '-w', filename]
    objdumpProc = Popen(objdumpCmd, stdout=PIPE)

    objdumpOutput = objdumpProc.communicate()[0]
    if objdumpProc.returncode != 0:
        logging.error('Could not dump %s' % filename)
        sys.exit(-1)

    for line in [l.decode() for l in objdumpOutput.splitlines()] :
        fields = line.split()
        if len(fields) <= 7:
            continue
        if fields[1] != sectionName:
            continue
        try:
            idx = int(fields[0])
            size = int(fields[2], 16)
            offset = int(fields[5], 16)
            return (size, offset)
        except ValueError:
            continue

    # The needed section could not be found
    raise Exception('Could not find "{0}" ELF section in "{1}"'.format(
                    sectionName,
                    filename)
                   )

# Read the entire content of an ELF section into a string
def getSectionContent(size, offset, filename):
    with open(filename, mode='rb') as f:
        f.seek(offset)
        d = ''
        try:
            c = f.read(size)
            d = c.decode()
        except UnicodeDecodeError:
            logging.error('Failed to read section containing:')
            print(c)
            raise
        # The linker pads sections with null bytes; our real data
        # cannot have null bytes because it is just text.  Discard
        # nulls.
        return d.replace('\0', '')



def handleExecutable(inputFile, llvmLinker, outputFile=None):
    (sectionSize, sectionOffset) = getSectionSizeAndOffset(elfSectionName, inputFile)
    if sectionSize == 0:
        logging.error('%s is empty' % elfSectionName)
        return 1

    content = getSectionContent(sectionSize, sectionOffset, inputFile)
    fileNames = content.split('\n')

    if outputFile == None:
      outputFile = inputFile + '.' + moduleExtension

    logging.debug('Found the following file names from ELF header:\n' + pprint.pformat(fileNames))
    linkCmd = [ llvmLinker, '-o', outputFile ]
    linkCmd.extend([x for x in fileNames if x != ''])
    logging.info('Writing output to {0}'.format(outputFile))
    linkProc = Popen(linkCmd)
    exitCode = linkProc.wait()
    return exitCode

def handleArchive(inputFile, llvmArchiver, outputFile=None):
    inputFile = os.path.abspath(inputFile)
    originalDir = os.getcwd() # This will be the destination
    
    # Make temporary directory to extract objects to
    tempDir = ''
    bitCodeFiles = [ ]
    retCode=0
    try:
        tempDir = tempfile.mkdtemp(suffix='wllvm')
        os.chdir(tempDir)

        # Extract objects from archive
        arC = ['ar','x',inputFile]
        arP = Popen(arC)
        arPE = arP.wait()

        if arPE != 0:
            errorMsg = 'Failed to execute archiver with command {0}'.format(arC)
            logging.error(errorMsg)
            raise Exception(errorMsg)

        # Iterate over objects and examine their .llvm_bc header
        for (root, dirs, files) in os.walk(tempDir):
           logging.debug('Exploring "{0}"'.format(root))
           for f in files:
               fPath = os.path.join(root, f)
               if FileType.getFileType(fPath) == FileType.OBJECT:
                   # Extract bitcode locations from object
                   (sectionSize, sectionOffset) = getSectionSizeAndOffset(elfSectionName, fPath)
                   content = getSectionContent(sectionSize, sectionOffset, fPath)
                   fileNames = content.split('\n')
                   for bcFile in fileNames:
                       if bcFile != '':
                           if not os.path.exists(bcFile):
                               logging.warning('{0} ELF section in {1} lists bitcode library "{2}" but it could not be found'.format(
                                               elfSectionName, f, bcFile))
                           else:
                               bitCodeFiles.append(bcFile)
               else:
                   logging.warning('Ignoring file "{0}" in archive'.format(f))
        
        logging.info('Found the following bitcode file names to build bitcode archive:\n{0}'.format(
                      pprint.pformat(bitCodeFiles)))

    finally:
        # Delete the temporary folder
        logging.debug('Deleting temporary folder "{0}"'.format(tempDir))
        shutil.rmtree(tempDir)

    # Build bitcode archive
    os.chdir(originalDir)

    # Pick output file path if outputFile not set
    if outputFile == None:
      if inputFile.endswith('.a'):
        # Strip off .a suffix
        outputFile = inputFile[:-2]
      else:
        outputFile = inputFile
      outputFile +='.' + bitCodeArchiveExtension

    logging.info('Writing output to {0}'.format(outputFile))

    # We do not want full paths in the archive so we need to chdir into each
    # bitcode's folder. Handle this by calling llvm-ar once for all bitcode
    # files in the same directory

    # Map of directory names to list of bitcode files in that directory
    dirToBCMap = {}
    for bitCodeFile in bitCodeFiles:
        dirName = os.path.dirname(bitCodeFile)
        basename = os.path.basename(bitCodeFile)
        if dirName in dirToBCMap:
            dirToBCMap[dirName].append(basename)
        else:
            dirToBCMap[dirName] = [ basename ]

    logging.debug('Built up directory to bitcode file list map:\n{0}'.format(
                   pprint.pformat(dirToBCMap)))

    for (dirname, bcList) in dirToBCMap.items():
        logging.debug('Changing directory to "{0}"'.format(dirname))
        os.chdir(dirname)
        larCmd = [llvmArchiver, 'rs', outputFile ] + bcList
        larProc = Popen(larCmd)
        retCode = larProc.wait()
        if retCode != 0:
          logging.error('Failed to execute:\n{0}'.format(pprint.pformat(larCmd)))
          break


    if retCode == 0:
        logging.info('Generated LLVM bitcode archive {0}'.format(outputFile))
    else:
        logging.error('Failed to generate LLVM bitcode archive')

    return retCode

def main(args):
    import argparse

    llvmToolPrefix = os.getenv(llvmCompilerPathEnv)
    if not llvmToolPrefix:
      llvmToolPrefix = ''
    llvmLinker = os.path.join(llvmToolPrefix, 'llvm-link')
    llvmArchiver = os.path.join(llvmToolPrefix, 'llvm-ar')

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("wllvm_binary", help="A binary produced by wllvm/wllvm++")
    parser.add_argument("--linker","-l", 
                        help='The LLVM bitcode linker to use. Default "%(default)s"',
                        default=llvmLinker)
    parser.add_argument("--archiver","-a", 
                        help='The LLVM bitcode archiver to use. Default "%(default)s"',
                        default=llvmArchiver)
    parser.add_argument("--output","-o", 
                        help='The output file. Defaults to a file in the same directory ' + 
                             'as the input with the same name as the input but with an ' +
                             'added file extension (.'+ moduleExtension + ' for bitcode '+
                             'modules and .' + bitCodeArchiveExtension +' for bitcode archives)',
                        default=None)
    parsedArgs = parser.parse_args()

    inputFile = parsedArgs.wllvm_binary
    llvmLinker= parsedArgs.linker

    # Check file exists
    if not os.path.exists(inputFile):
        logging.error('File "{0}" does not exist.'.format(inputFile))
        return 1

    # Check output destitionation if set
    outputFile = parsedArgs.output
    if outputFile != None:
      # Get Absolute output path
      outputFile = os.path.abspath(outputFile)
      if not os.path.exists(os.path.dirname(outputFile)):
        logging.error('Output directory "{0}" does not exist.'.format(
                      os.path.dirname(outputFile)))
        return 1
      
    ft = FileType.getFileType(inputFile)
    logging.debug('Detected file type is {0}'.format(FileType.revMap[ft]))

    if ft == FileType.EXECUTABLE:
        logging.info('Generating LLVM Bitcode module')
        return handleExecutable(inputFile, llvmLinker, outputFile )
    elif ft == FileType.ARCHIVE:
        logging.info('Generating LLVM Bitcode archive')
        return handleArchive(inputFile, llvmArchiver, outputFile )
    else:
        logging.error('File "{0}" of type {1} cannot be used'.format(inputFile,FileType.revMap[ft]))
        return 1
    sys.exit(exitCode)

if __name__ == '__main__':
    sys.exit(main(sys.argv))