url_samples/urldbsample_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>
 
#ifdef WIN32
#       include <winsock2.h>
#endif
 
#include "dca/dca_base.h"
#include "dca/dca_url_classification.h"
#include "dca/dca_callbacks.h"
 
#include "mythreads.h"
#include "schedule_event_subscriber.h"
 
using namespace dca;
 
const std::string S_ToolName = "urldbsample_extended";
const std::string S_ToolVersion = "1.3";
 
const std::string S_UsageString =
        " <redist-folder> <ticket> <product> <url-list-file>\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\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& redist_folder, InitData& initData )
{
        initData.binDir = redist_folder + DCA_BINDIR;
        initData.initDir = redist_folder + DCA_INITDIR;
        initData.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& 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 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" );
 
                if( myCategory != NullCategory ) {
                        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 CategoriesInfo& myCategoriesInfo )
{
        std::cout << "Starting URL db classification routine..." <<  std::endl;
 
        UrlClassificationResults myUrlClassificationResults;
 
        std::vector<std::string> myUrlList;
        LoadUrlFile( aUrlListFile, myUrlList );
 
        for( std::vector<std::string>::const_iterator U = myUrlList.begin(),
                UEnd = myUrlList.end(); U != UEnd; ++U ) {
 
                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
                FunctionResult myFR = myUrlDbClassifier.classify( myUrl, myUrlClassificationResults );
 
                if( !myFR ) {
                        // error occured.
                        std::cerr << "Error from uURLrl 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;
}
 
int main( int argc, char *argv[] )
{
        PrintToolHeader();
 
        int rc = 0;
 
        try {
 
                if( argc < 5 ) {
                        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 ];
 
                if( myRedistFolder.empty() || myTicket.empty() ||
                         myProduct.empty() || myUrlList.empty() ) {
                        PrintUsage( argv[0] );
                        return 5;
                }
 
                // check for traling fileslash - and add if necessary
                const char c = myRedistFolder[ myRedistFolder.length() - 1 ];
                if( c != '/' && c != '\\' )
                        myRedistFolder += "/";
 
                // Initialize socket on Windows and 3rd party libraries
                if( !StartupLibraries() ) 
                        return 5;
 
                {
                        // 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 );
 
                                // initialize the URL classification module and create a URL db classifier
                                // (embedded URL detection enabled, statistics and unknown url upload enabled)
                                const UrlClassification myUrlClassification = UrlClassification::create( myDca, myLicense );
 
                                UrlDbClassifierOptions myUrlDbClassifierOptions;
                                myUrlDbClassifierOptions.enable_EmbeddedUrlDetection = true;
                                myUrlDbClassifierOptions.detect_EmbeddedUrlsInUrlPath = true;
                                myUrlDbClassifierOptions.enable_Feedback = true;
 
                                const UrlDbClassifier myUrlDbClassifier = myUrlClassification.createDbClassifier( myDbConnection, myUrlDbClassifierOptions );
 
                                // create a categories info for printing out the category 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 );
 
                                // set up an instance of MyScheduleEventSubscriber to capture the schedule event information
                                MyScheduleEventSubscriber mySubscriber;
 
                                StartupThreads(myDca, myUpdateModule, &mySubscriber);
 
                                // call URL db classification routine
                                TestUrlClassification( myUrlList, myDca, myUrlDbClassifier, myCategoriesInfo );
 
                                myUpdateModule.cancelUpdate();
                                myDca.signal( DCA_SIG_ABORT );
                                ShutdownThreads();
 
                                mySubscriber.DumpEventMessages();
                        }
                }
        }
        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;
}