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
bitcode section from the provided ELF or MACH-O object and reassemble
it into an actual bitcode file.  The ELF or MACH-Osection 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 or MACH-O objects and read their bitcode sections and
create a LLVM Bitcode archive from the bitcode files.

The above language is deliberately vague, since ELF contains a
.llvm_bc section, whereas the MACH-O contains a segment called __LLVM
that contains a section called __llvm_bc.

"""

import os
import sys
import subprocess as sp
#from subprocess import *
from driver.utils import llvmCompilerPathEnv
from driver.popenwrapper import Popen
from driver.utils import elfSectionName
from driver.utils import darwinSegmentName
from driver.utils import darwinSectionName
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'

#verbose flag to pass to the external calls
verboseFlag = False

# 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=sp.PIPE)

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

    for line in [l.decode('utf-8') 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('utf-8')
        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', '')

# os independent abstraction (darwin version)
#use otool to extract the segment and section
#iam: "xxd -r" would just refuse to play nice, so it got tossed.
# if any one can get
#    otool -X -s __LLVM __llvm_bc inputFile | xxd -r
# to work this can be simplified.
def extract_section_darwin(inputFile):

    otoolCmd  = ['otool', '-X', '-s', darwinSegmentName, darwinSectionName, inputFile]
    otoolProc = Popen(otoolCmd, stdout=sp.PIPE)

    otoolOutput = otoolProc.communicate()[0]
    if otoolProc.returncode != 0:
        logging.error('otool failed on %s' % inputFile)
        sys.exit(-1)

    lines = otoolOutput.splitlines()
    octets = []
    for line in lines:
        (_, octetline) = line.split('\t')
        octets.extend(octetline.split())
    octets = ''.join(octets)
    contents = octets.decode('hex').splitlines()
    if not contents:
        logging.error('{0} contained no {1} segment'.format(inputFile, darwinSegmentName))
    return contents

# os independent abstraction (*nix version)
#analagous procedure for linux (relying on getSectionSizeAndOffset and getSectionContent)
def extract_section_linux(inputFile):
    (sectionSize, sectionOffset) = getSectionSizeAndOffset(elfSectionName, inputFile)
    content = getSectionContent(sectionSize, sectionOffset, inputFile)
    contents = content.split('\n')
    if not contents:
        logging.error('{0} contained no {1} section is empty'.format(inputFile, elfSectionName))
    return contents


def handleExecutable(inputFile, outputFile, extractor, llvmLinker):

    fileNames = extractor(inputFile)

    if not fileNames:
        return 1

    if outputFile == None:
        outputFile = inputFile + '.' + moduleExtension
        
    linkCmd = [ llvmLinker, '-o', outputFile ]
    
    if verboseFlag:
        linkCmd.insert(1, '-v')
    
    linkCmd.extend([x for x in fileNames if x != ''])
    logging.info('Writing output to {0}'.format(outputFile))
    linkProc = Popen(linkCmd)
    exitCode = linkProc.wait()
    logging.info('{0} returned {1}'.format(llvmLinker, str(exitCode)))
    return exitCode




def handleArchive(inputFile, outputFile, arCmd, fileType, extractor, llvmArchiver):
    inputFile = os.path.abspath(inputFile)
    originalDir = os.getcwd() # This will be the destination

    arCmd.append(inputFile);
    if verboseFlag:
        arCmd.insert(1, '-v')
    
    # Make temporary directory to extract objects to
    tempDir = ''
    bitCodeFiles = [ ]
    retCode=0
    try:
        tempDir = tempfile.mkdtemp(suffix='wllvm')
        os.chdir(tempDir)

        # Extract objects from archive
        arP = Popen(arCmd)
        arPE = arP.wait()

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

        # Iterate over objects and examine their bitcode inserts
        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:

                   # Extract bitcode locations from object
                   contents = extractor(fPath)

                   for bcFile in contents:
                       if bcFile != '':
                           if not os.path.exists(bcFile):
                               logging.warning('{0} lists bitcode library "{1}" but it could not be found'.format(f, bcFile))
                           else:
                               bitCodeFiles.append(bcFile)
               else:
                   logging.info('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)

    return buildArchive(inputFile, outputFile, llvmArchiver, bitCodeFiles)

def buildArchive(inputFile, outputFile, llvmArchiver, bitCodeFiles):
    retCode=0
    # 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 extract_bc_args(args):
    import argparse

    global verboseFlag
    
    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("--verbose","-v",
                        help='Call the external procedures in verbose mode.',
                        action="store_true")
    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
    verboseFlag = parsedArgs.verbose
    
    # Check file exists
    if not os.path.exists(inputFile):
        logging.error('File "{0}" does not exist.'.format(inputFile))
        return (False, inputFile, '', llvmLinker, llvmArchiver)

    # 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 (False, inputFile, '', llvmLinker, llvmArchiver)

    return (True, inputFile, outputFile, llvmLinker, llvmArchiver)
    

def main(args):

    (success, inputFile, outputFile, llvmLinker, llvmArchiver) = extract_bc_args(args)

    if not success:
        return 1

    if ( sys.platform.startswith('freebsd') or
         sys.platform.startswith('linux') ):

        process_file_unix(inputFile, outputFile, llvmLinker, llvmArchiver)

    elif sys.platform.startswith('darwin'):

        process_file_darwin(inputFile, outputFile, llvmLinker, llvmArchiver)

    else:
        #iam: do we work on anything else? 
        logging.error('Unsupported or unrecognized platform: {0}'.format(sys.platform))
        return 1

    
def process_file_unix(inputFile, outputFile, llvmLinker, llvmArchiver):
    ft = FileType.getFileType(inputFile)
    logging.debug('Detected file type is {0}'.format(FileType.revMap[ft]))

    extractor = extract_section_linux
    arCmd     = ['ar','x']
    ofileType = FileType.ELF_OBJECT

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

    

def process_file_darwin(inputFile, outputFile, llvmLinker, llvmArchiver):
    ft = FileType.getFileType(inputFile)
    logging.debug('Detected file type is {0}'.format(FileType.revMap[ft]))

    extractor = extract_section_darwin
    arCmd     = ['ar','-x']
    ofileType = FileType.MACH_OBJECT

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


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