customdb_samples/customdbsample_extended/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 "dca/dca_base.h"
#include "dca/dca_url_classification.h"
#include "dca/dca_callbacks.h"
 
#include "mythreads.h"
 
using namespace dca;
 
const std::string S_ToolName = "customdbsample_extended";
const std::string S_ToolVersion = "1.2";
 
const std::string S_UsageString =
        " <redist-folder> <ticket> <product> <url-list-file> <custom-db-folder>\n"
        "  redist-folder - the folder where the DCA is installed to\n"
        "  ticket - a valid ticket\n"
        "  product - the product associated with your ticket\n"
        "  url-list-file - file that includes the URLs to classify\n"
        "  custom-db-folder - folder where the Custom Database is located\n\n"
        ;
 
volatile bool S_schedulerThreadError = false;
volatile bool S_updateThreadError = false;
 
#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& redist_folder, InitData& initData )
{
        initData.binDir = redist_folder + DCA_BINDIR;
        initData.initDir = redist_folder + DCA_INITDIR;
        initData.logDir = DCA_LOGDIR;
}
 
static void SetupLicense( const std::string& ticket, const std::string& product,
        LicenseData& licenseData )
{
        licenseData.ticket      = ticket;
        licenseData.product     = product;
}
 
static void SetupConnectionData( DbConnectionData& cData )
{
        cData.useLocalDatabase = true;
        cData.dbType = DBT_Url;
}
 
static void SetupCustomConnectionData( const std::string& customdb_folder, DbConnectionData& cData )
{
        cData.useLocalDatabase = true;
        cData.dbType = DBT_Custom;
        cData.customData.configDir = customdb_folder;
}
 
static void PrintResults( const CategoriesInfo& catinfos, const UrlClassificationResults& cats )
{
        const DCA_SIZE_TYPE numOfCats = cats.size();
 
        if( numOfCats == 0 ) {
                std::cout << "No categories found." << std::endl;
                return;
        }
 
        const Categories myCategories = catinfos.getCategories();
 
        for( DCA_INDEX_TYPE i = 0; i < numOfCats; i++ ) {
 
                const UrlClassificationResult result = cats[ i ];
                const Category myCategory = myCategories.byId( result );
                const std::string catname = myCategory.name( "en_US" );
 
                std::cout << (i+1) << ".\t Category '" << catname <<
                        "' (id=" << myCategory.id() <<
                        ", groupid=" << myCategory.groupId() <<
                        ")" << std::endl;
        }
}
 
static void PrintToolHeader()
{
        std::cout << "IBM DCA Sample: " << S_ToolName << " (" << S_ToolVersion << ")" << std::endl;
}
 
static void PrintUsage( const char *name )
{
        std::cout << "usage: " << name << S_UsageString << std::endl;
}
 
static void LoadUrlFile( const std::string& fileName, std::vector<std::string>& urlList )
{
        std::ifstream fstream( fileName.c_str(), std::ios::in );
        if (!fstream.is_open()) return;
 
        std::string line;
 
        while ( std::getline(fstream, line) )
        {
                if( !line.empty() && line[line.length()-1] == '\r')
                        line.erase( line.length() - 1 );
 
                if( !line.empty() )
                        urlList.push_back( line );
        }
}
 
void TestUrlClassification( const std::string& aUrlListFile, const DcaInstance& myDca,
        const UrlDbClassifier& myUrlDbClassifier, const UrlDbClassifier& myUrlCustomDbClassifier,
        const CategoriesInfo& myCategoriesInfo )
{
        try {
                std::cout << "Starting custom & URL db classification routine..." <<  std::endl;
 
                UrlClassificationResults myUrlClassificationResults;
 
                std::vector<std::string> myUrlList;
                LoadUrlFile( aUrlListFile, myUrlList );
 
                // for each found URL first check for a match in the Custom Database
                //  -- if matched - done
                //  -- else try to match given URL in standard URL database
                for( std::vector<std::string>::const_iterator U = myUrlList.begin(),
                        UEnd = myUrlList.end(); U != UEnd; ++U ) {
 
                        // check for errors from update and scheduler thread....
                        if( S_schedulerThreadError ) {
                                std::cerr << "Scheduler Thread got an error from DcaInstance::schedule() call. aborting." <<
                                        std::endl;
                                return;
                        }
 
                        if( S_updateThreadError ) {
                                std::cerr << "Updater Thread got an error from UpdateModule::performUpdate() call. aborting." <<
                                        std::endl;
                                return;
                        }
 
                        // start over with next URL
                        const std::string& myUrlString = *U;
 
                        std::cout << "  Starting URL db classification for URL '" <<
                                myUrlString << "'" << std::endl;
 
                        // setup a URL for given string
                        const Url myUrl  = Url::create( myDca, myUrlString );
 
                        // start the classification
                        // match with Custom Database
                        FunctionResult myFR = myUrlCustomDbClassifier.classify( myUrl, myUrlClassificationResults );
                        if( !myFR ) {
                                std::cerr << "Error from custom URL db classification. Details: " << myFR.getDescription() <<
                                        " (" << myFR.getReturnCode() << "). Continuing with next URL." << std::endl;
                                continue;
                        }
 
                        if( !myUrlClassificationResults.isUnknownUrl() ) {
                                std::cout << "Got a match from Custom Database." << std::endl;
 
                                if( !myUrlClassificationResults.isCategorized() ) {
                                        std::cout << "Results: URL '" << myUrlString << "' is not categorized (but known in database). "
                                                << "Continuing with next URL" << std::endl;
                                }
                                else {
                                        std::cout << "Classification Results for URL '" << myUrlString << "'" << std::endl;
                                        PrintResults( myCategoriesInfo, myUrlClassificationResults );
                                }
                                continue;
                        }
 
                        // match with local URL database
                        myFR = myUrlDbClassifier.classify( myUrl, myUrlClassificationResults );
 
                        if( !myFR ) {
                                std::cerr << "Error from URL db classification. Details: " << myFR.getDescription() <<
                                        " (" << myFR.getReturnCode() << "). Continuing with next URL." << std::endl;
                                continue;
                        }
 
                        if( myUrlClassificationResults.isUnknownUrl() ) {
                                // the URL is not known in the database
                                std::cout << "Results: URL '" << myUrlString << "' is not known in the database. " <<
                                        "Continuing with next URL." << std::endl;
                                continue;
                        }
 
                        if( !myUrlClassificationResults.isCategorized() ) {
                                // the URL is known but does not contains any category. This is either a white-host or
                                // includes only categorized sub-folders.
                                std::cout << "Results: URL '" << myUrlString << "' is not categorized (but known in database). "
                                        << "Continuing with next URL" << std::endl;
                                continue;
                        }
 
                        // given URL is known and returned matched categories
                        std::cout << "Classification Results for URL '" << myUrlString << "'" << std::endl;
                        PrintResults( myCategoriesInfo, myUrlClassificationResults );
                }
                std::cout << "Leaving URL db classification routine." << std::endl;
        }
        catch( const ExDca& ex ) {
                std::cerr << "DCA Exception occured in TestUrlClassification(). Details: " <<
                        ex.getDescription() << " (" << ex.getReturnCode() << ")." << std::endl;
        }
        catch( const std::exception& s ) {
                std::cerr << "std::exception occured in TestUrlClassification(). Details: " <<
                        s.what() << "." << std::endl;
        }
        catch(...) {
                std::cerr << "Unknown exception caught in TestUrlClassification()." << std::endl;
        }
}
 
int main( int argc, char *argv[] )
{
        PrintToolHeader();
 
        int rc = 0;
 
        try {
 
                if( argc != 6 ) {
                        PrintUsage( argv[0] );
                        return 5;
                }
 
                std::string myRedistFolder = argv[ 1 ];
                const std::string myTicket = argv[ 2 ];
                const std::string myProduct = argv[ 3 ];
                const std::string myUrlList = argv[ 4 ];
                std::string myDbFolder = argv[ 5 ];
 
                if( myRedistFolder.empty() || myTicket.empty() ||
                         myProduct.empty() || myUrlList.empty()||
                         myDbFolder.empty() ) {
                        PrintUsage( argv[0] );
                        return 5;
                }
 
                // check for trailing fileslash - and add if necessary
                char c = myRedistFolder[ myRedistFolder.length() - 1 ];
                if( c != '/' && c != '\\' )
                        myRedistFolder += "/";
 
                c = myDbFolder[ myDbFolder.length() - 1 ];
                if( c != '/' && c != '\\' )
                        myDbFolder += "/";
 
                // init the 3rd party libraries
                InitCUrl();
                SetOpenSslCallbacks();
 
                {
                        // setup DCA directories
                        InitData myInitData;
                        SetupInitData( myRedistFolder, myInitData );
 
                        // instantiate DCA API
                        DcaInstance myDca;
                        myDca = DcaInstance::create( myInitData );
 
                        // setup license data
                        LicenseData myLicenseData;
                        SetupLicense( myTicket, myProduct, myLicenseData );
                        const License myLicense = myDca.createLicense( myLicenseData );
 
                        if( !myLicense.isLicensed( UrlClassification::ID ) ) {
                                std::cout << "DCA is not licensed!" << std::endl;
                                rc = 5;
                        } else {
                                // setup a signature database connection
                                DbConnectionData myDbConnectionData;
                                SetupConnectionData( myDbConnectionData );
                                const DbConnection myDbConnection = myDca.createDbConnection( myLicense, myDbConnectionData );
 
                                // setup a custom db signature database connection
                                DbConnectionData myCustomDbConnectionData;
                                SetupCustomConnectionData( myDbFolder, myCustomDbConnectionData );
                                const DbConnection myCustomDbConnection = myDca.createDbConnection( myLicense, myCustomDbConnectionData );
 
                                // initialize the URL classification module and create a URL db classifier
                                UrlClassification myUrlClassification = UrlClassification::create( myDca, myLicense );
                                const UrlDbClassifier myUrlDbClassifier = myUrlClassification.createDbClassifier( myDbConnection );
 
                                // initialize the URL classifier connected to the custom DbConnection
                                const UrlDbClassifier myUrlCustomDbClassifier = myUrlClassification.createDbClassifier( myCustomDbConnection );
 
                                // create a categories info for printing out the cotegory names together with the results per URL
                                const CategoriesInfo myCategoriesInfo = myDca.getCategoriesInfo( DCA_CAT_INFO_TYPE_URL );
 
                                // create the update module
                                const UpdateModule myUpdateModule = UpdateModule::create( myDca, myLicense );
 
                                StartupThreads(myDca, myUpdateModule, &S_schedulerThreadError, &S_updateThreadError);
 
                                // call URL Classification with a Custom Database
                                TestUrlClassification( myUrlList, myDca, myUrlDbClassifier, myUrlCustomDbClassifier, myCategoriesInfo );
 
                                myUpdateModule.cancelUpdate();
                                myDca.signal( DCA_SIG_ABORT );
                                ShutdownThreads();
                        }
                }
        }
        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
        UnsetOpenSslCallbacks();
        DeinitCUrl();
 
        return rc;
}