malware_samples/malwaresample/main.cpp

/* IBM Source Code */
/* (C) Copyright IBM Corp. 2009, 2012 */
/* Licensed Materials - Property of IBM */
/* US Government Users Restricted Rights - Use duplication or disclosure restricted by GSA Schedule Contract with IBM Corp. */
 
#include <string>
#include <vector>
#include <iostream>
#include <fstream>
#include <ctime>
 
#include "dca/dca_base.h"
#include "dca/dca_malware_classification.h"
#include "dca/dca_callbacks.h"
 
#ifdef WIN32
#       include <winsock2.h>
#endif
 
using namespace dca;
using namespace dca_malware;
 
const std::string S_ToolName = "malwaresample";
const std::string S_ToolVersion = "1.0";
 
const std::string S_UsageString =
        "<dca-redist-folder> <ticket> <product> <input-file>\n"
        "  dca-redist-folder - the folder where the DCA is installed\n"
        "  ticket - a valid ticket\n"
        "  product - the product associated with your ticket\n"
        "  input-file - a file with sample hashes\n\n"
        ;
 
#ifdef WIN32
#       define DCA_BINDIR               "bin/Win32"
#else
#       define DCA_BINDIR               "bin/linux"
#endif
 
#define DCA_INITDIR             "init"
 
#define DCA_LOGDIR              "./logs"
 
static void SetupInitData( const std::string& strRedistFolder,
                                                   InitData& aInitData )
{
        aInitData.binDir        = strRedistFolder + DCA_BINDIR;
        aInitData.initDir       = strRedistFolder + DCA_INITDIR;
        aInitData.logDir        = DCA_LOGDIR;
}
 
static bool StartupLibraries()
{
#ifdef WIN32
        // Windows needs an extra socket-startup for this process to work 
        // correctly with e.g. IP(v6) input IP addresses
        WORD wVersionRequested = MAKEWORD( 2, 2 );
        WSADATA wsaData;
        int err = WSAStartup( wVersionRequested, &wsaData );
        if ( err != 0 ) {
                std::cout << "Error on WSAStartup (" << err << ") occured, aborting" << 
                        std::endl;
                        return false;
        }
#endif  
 
        // init the 3rd party libraries
        InitCUrl();
        SetOpenSslCallbacks();
        return true;
}
 
static void ShutdownLibraries()
{
        // deinit the 3rd party libraries
        UnsetOpenSslCallbacks();
        DeinitCUrl();
        
#ifdef WIN32    
        // Cleanup Windows sockets for this process
        WSACleanup();
#endif
}
 
static void SetupLicense( const std::string& strTicket,
                                                  const std::string& strProduct,
                                                  LicenseData& aLicenseData )
{
        aLicenseData.ticket             = strTicket;
        aLicenseData.product    = strProduct;
}
 
static void SetupConnectionData( DbConnectionData& aDbConnectionData )
{
        aDbConnectionData.useLocalDatabase      = true;
        aDbConnectionData.dbType                        = DBT_Malware;
}
 
static void PrintDbConnectionInfo( const DbConnection& aDbConnection )
{
        DatabaseInformation databaseInformation = 
                aDbConnection.getDatabaseInformation();
 
        std::cout << "Malware Database Version: " << databaseInformation.versionString 
                << " as of " << databaseInformation.creationDateUTC << std::endl;
}
 
static void PrintLicenseInfo( const License& aLicense )
{
        const time_t expirationDate = aLicense.getExpirationDate();
        struct tm *expirationTime = localtime( &expirationDate );
 
        std::cout << "License Info:" << std::endl;
        std::cout << "  DCA is " << ( aLicense.isLicensed() ? "licensed." : 
                "not licensed." ) << std::endl;
        std::cout << "  MaxUsers:" << aLicense.getMaxUsers() << 
                std::endl;
        std::cout << "  MaxSessions:" << aLicense.getMaxSessions() << 
                std::endl;
        std::cout << "  Ticket:" << aLicense.getTicket() << 
                std::endl;
        std::cout << "  Session:" << aLicense.getSession() << 
                std::endl;
        std::cout << "  Last Message:" << aLicense.getLastMessage() << 
                std::endl;
        std::cout << "  Expiration Date:" << asctime( expirationTime ) << 
                std::endl;
}
 
static void PrintToolHeader()
{
        std::cout << "IBM DCA Sample: " << S_ToolName << " (" << S_ToolVersion << 
                ")" << std::endl;
}
 
static void PrintUsage( const char *pSampleName )
{
        std::cout << pSampleName << " usage:" << std::endl;
        std::cout << S_UsageString << std::endl;
}
 
void PrintResult( const MalwareClassificationResult& aMalwareClassificationResult,
                                  const MalwareCategories& aMalwareCategories,
                                  const MalwareEnums& aMalwareEnums,
                                  const std::string& strIndent = "",
                                  const std::string& strValueDesc = "" )
{
        // each result item contains of a category id and a value related to the
        // category
        DCA_CATEGORY_ID_TYPE myCatId = aMalwareClassificationResult.categoryId;
        int myValue = aMalwareClassificationResult.value;
 
        // retrieve the category object to print out its name
        MalwareCategory myMalwareCategory = aMalwareCategories.byId( myCatId );
 
        std::cout << strIndent << "MalwareCategory '" << myMalwareCategory.name() << 
                "' (id=" << myCatId << ")" << std::endl;
 
        // if the matched category contains one or more related MalwareEnum objects,
        // print out the matched one
        const DCA_SIZE_TYPE numOfEnums = myMalwareCategory.enumsSize();
        if( numOfEnums > 0 ) {
                // the retrieved value points to a unique MalwareEnum object id
                int enumId = myMalwareCategory.enumByValue( myValue );
                // retrieve the object with the enumId
                const MalwareEnum myMalwareEnum = aMalwareEnums.byId( enumId );
                std::cout << strIndent << "\t MalwareEnum '" << myMalwareEnum.name() << "' (id=" <<
                        enumId << ")" << std::endl;
        }
        else {
                // no Enum object retrieved print out the value
                std::cout << strIndent << "\tValue " << myValue << strValueDesc << std::endl;
        }
}
 
void PrintResults( const MalwareClassificationResults& aMalwareClassificationResults,
                                   const MalwareCategoriesInfo& aMalwareCategoriesInfo )
{
        if( aMalwareClassificationResults.isUnknownMalware() ) {
                std::cout << "Results: Unknown" << std::endl;
                return;
        }
        if( !aMalwareClassificationResults.isCategorized() ) {
                std::cout << "Results: None" << std::endl;
                return;
        }
 
        const DCA_SIZE_TYPE numOfResults = aMalwareClassificationResults.size();
 
        std::cout << "Results (" << numOfResults << ") categories" << std::endl;
 
        const MalwareCategories myMalwareCategories = aMalwareCategoriesInfo.getCategories();
        const MalwareEnums myMalwareEnums = aMalwareCategoriesInfo.getEnums();
 
        for( DCA_INDEX_TYPE i = 0; i < numOfResults; ++i ) {
                const MalwareClassificationResult myMalwareClassificationResult = aMalwareClassificationResults[ i ];
                PrintResult( myMalwareClassificationResult, myMalwareCategories, myMalwareEnums );
        }
}
 
void DoCallMalwareClassification( const std::vector<std::string>& aHashStrings,
                                                          const MalwareDbClassifier& aMalwareClassifier,
                                                          const MalwareCategoriesInfo& aMalwareCategoriesInfo )
{
        std::vector< std::string >::const_iterator I = aHashStrings.begin();
        std::vector< std::string >::const_iterator IEnd = aHashStrings.end();
 
        MalwareClassificationResults myMalwareClassificationResults;
        FunctionResult myFR;
 
        for( ; I != IEnd; ++I ) {
                const std::string& strHash = *I;
 
                std::cout << std::endl << "MALWARE Classification ('" << strHash << "')..." << std::endl;
                unsigned char hashBuffer[33] = {
                        0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0,
                        0
                };
                const char* hashPtr = strHash.c_str();
                bool data_ok = true;
                if (strHash.length() >= 64 && hashPtr) {
                        // hash string
                        for (int i = 0; i < 32; i++) {
                                unsigned char value = 0;
                                char char0 = *(hashPtr + (2*i));
                                char char1 = *(hashPtr + (2*i) + 1);
                                if (char0 >= 'A' && char0 <= 'F') {
                                        char0 = char0 + 32;
                                }
                                if (char0 >= '0' && char0 <= '9') {
                                        value = (unsigned char)((char0 - '0') << 4);
                                }
                                else if (char0 >= 'a' && char0 <= 'f') {
                                        value = (unsigned char)((char0 - 'A' + 10) << 4);
                                }
                                else {
                                        data_ok = false;
                                }
                                if (char1 >= 'A' && char1 <= 'F') {
                                        char1 = char1 + 32;
                                }
                                if (char1 >= '0' && char1 <= '9') {
                                        value |= ( (char1 - '0') & 0x0fU );
                                }
                                else if (char1 >= 'a' && char1 <= 'f') {
                                        value |= ( (char1 - 'A' + 10) & 0x0fU );
                                }
                                else {
                                        data_ok = false;
                                }
                                if (data_ok) {
                                        hashBuffer[i] = value;
                                }
                        }
                }
 
                if (data_ok) {
                        const Hash myHash(hashBuffer, 32);
 
                        myFR = aMalwareClassifier.classify( myHash, myMalwareClassificationResults );
 
                        if( !myFR ) {
                                std::cout << "Error: rc=" << myFR.getReturnCode() <<
                                        ", description=" << myFR.getDescription() << std::endl;
                        }
                        else PrintResults( myMalwareClassificationResults, aMalwareCategoriesInfo );
                }
        }
}
 
void TestMalwareClassification( const MalwareDbClassifier& aMalwareClassifier,
                                                    const MalwareCategoriesInfo& aMalwareCategoriesInfo,
                                                        const std::string& strFileName )
{
        std::ifstream fstream( strFileName.c_str(), std::ios::in );
        if ( !fstream.is_open() ) {
                std::cout << "Error: File '" << strFileName << "' could not be opened!" << std::endl;
                return;
        }
 
        std::string strLine;
        std::vector< std::string > myHashStrings;
 
        while( std::getline( fstream, strLine ) ) {
                if( strLine.empty() )
                        continue;
 
                // delete all \r\n at the end of the line
                while( strLine[strLine.length()-1] == '\r' ||
                           strLine[strLine.length()-1] == '\n' )
                        strLine.erase( strLine.length() - 1 );
 
                if( strLine.empty() )
                        continue;
 
                if( strLine[0] == '#' ) // allow comment lines started with a # at first char
                        continue;
 
                // add line to myHashStrings
                myHashStrings.push_back( strLine );
        }
 
        fstream.close();
 
        if( !myHashStrings.empty() )
                DoCallMalwareClassification( myHashStrings, aMalwareClassifier, aMalwareCategoriesInfo );
        else
                std::cout << "Could not read any hash from file!" << std::endl;
}
 
int main( int argc, char *argv[] )
{
        PrintToolHeader();
 
        int rc = 5;
 
        try {
                if( argc < 5 ) {
                        PrintUsage( argv[0] );
                        return 5;
                }
 
                std::string strRedistFolder = argv[ 1 ];
                const std::string strTicket = argv[ 2 ];
                const std::string strProduct = argv[ 3 ];
                const std::string strFileName = argv[ 4 ];
 
                if( strRedistFolder.empty() ||
                        strTicket.empty() ||
                        strProduct.empty() ||
                        strFileName.empty() )
                {
                        PrintUsage( argv[0] );
                        return 5;
                }
 
                // check for trailing fileslash - and add if necessary
                const char c = strRedistFolder[ strRedistFolder.length() - 1 ];
                if( c != '/' && c != '\\' )
                        strRedistFolder += "/";
 
                // Initialize socket on Windows and 3rd party libraries
                if( !StartupLibraries() ) 
                        return 5;
 
                {
                        // setup DCA directories
                        InitData myInitData;
                        SetupInitData( strRedistFolder, myInitData );
 
                        // instantiate DCA API
                        DcaInstance myDca;
                        myDca = DcaInstance::create( myInitData );
 
                        // setup license data
                        LicenseData myLicenseData;
                        SetupLicense( strTicket, strProduct, myLicenseData );
                        const License myLicense = myDca.createLicense( myLicenseData );
 
                        PrintLicenseInfo( myLicense );
                        
                        if( myLicense.isLicensed( MalwareClassification::ID ) ) {
                                // initialize the Malware Classification module
                                const MalwareClassification myMalwareClassification = 
                                        MalwareClassification::create( myDca, myLicense );
 
                                DbConnectionData myDbConnectionData;
                                SetupConnectionData( myDbConnectionData );
 
                                // A DbConnection data to a local database must exist to use 
                                // MalwareClassification
                                const DbConnection myDbConnection = 
                                        myDca.createDbConnection( myLicense, myDbConnectionData );
                                PrintDbConnectionInfo( myDbConnection );
 
                                // create a MalwareDbClassifier to use in the TestMalwareClassification() 
                                // function
                                MalwareDbClassifierOptions myMalwareDbClassifierOptions;
                                myMalwareDbClassifierOptions.enable_Feedback = false;
                                const MalwareDbClassifier myClassifier = 
                                        myMalwareClassification.createDbClassifier( myDbConnection, myMalwareDbClassifierOptions );
 
                                // create a categories info for printing out the result 
                                // categories and values per input data
                                const MalwareCategoriesInfo myCategoriesInfo = 
                                        myMalwareClassification.getCategoriesInfo();
 
                                // call classification routine
                                TestMalwareClassification( myClassifier, myCategoriesInfo, 
                                        strFileName );
                                
                                rc = 0;
                        }
                }
        }
        catch( const ExDca& ex ) {
                std::cerr << "DCA Exception occured. Details: " <<  ex.getDescription() 
                        << " (" << ex.getReturnCode() << ")." << std::endl;
                rc = 10;
        }
        catch( const std::exception& s ) {
                std::cerr << "std::exception occured. Details: " << s.what() << "." << 
                        std::endl;
                rc = 10;
        }
        catch(...) {
                std::cerr << "Unknown exception caught." << std::endl;
                rc = 10;
        }
 
        // deinit the 3rd party libraries
        ShutdownLibraries();
 
        return rc;
}