// **********************************************************************
//
// Copyright (c) 1999
// Object Oriented Concepts, Inc.
// Billerica, MA, USA
//
// All Rights Reserved
//
// **********************************************************************

#include <OB/CORBA.h>
#include <OB/Util.h>
#include <OB/NTService.h>

//
// Static function to return a formatted error message for the last
// WIN32 error.
//
static char*
GetErrorText()
{
    DWORD rc;
    char* str = 0;
    char* cstr = 0;
    rc = ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER|
		       FORMAT_MESSAGE_FROM_SYSTEM|
		       FORMAT_MESSAGE_ARGUMENT_ARRAY,
		       0, GetLastError(), LANG_NEUTRAL, (char*)&str, 0, 0);
    if(str)
    {
	cstr = CORBA_string_dup(str);
	::LocalFree((HLOCAL)str);
    }
    else
    {
	cstr = CORBA_string_dup("FormatMessage failed");
    }
    return cstr;
}

//
// This is a custom message logger that is used to dump messages to
// the NT system log.
//
//
class OBNTEventViewer : public OBMessageViewer
{
    HANDLE eventSource_; // The event source
    CORBA_String_var service_; // Service name
    DWORD eventId_; // The event ID

    //
    // Report the event to the NT event log
    //
    void emitMessage(int, const char*);

    bool installValues(HKEY);

protected:

    OBNTEventViewer(const char*, DWORD);
    ~OBNTEventViewer();

public:

    static OBNTEventViewer* instance(const char*, DWORD);
      
    //
    // Install the registry entries for the event message resources
    //
    bool install();

    //
    // Remove the registry entries.
    //
    bool uninstall();
    
    //
    // Overrides for the MessageViewer methods
    //
    virtual void message(const char*);
    virtual void error(const char*);
    virtual void warning(const char*);
    virtual void nca(const char*);
};

// ----------------------------------------------------------------------
// OBNTEventViewer private and protected member implementation
// ----------------------------------------------------------------------

void
OBNTEventViewer::emitMessage(int eventType, const char* msg)
{
    const char* strings[1];
    
    strings[0] = msg;
    
    ::ReportEvent(eventSource_,         // handle of event source
		  eventType,            // event type
		  0,                    // event category
		  eventId_,             // event ID
		  NULL,                 // current user's SID
		  1,                    // strings in lpszStrings
		  0,                    // no bytes of raw data
		  strings,              // array of error strings
		  NULL);                // no raw data
}

bool
OBNTEventViewer::installValues(HKEY key)
{
    char path[512];
    
    //
    // Get path to executable.
    //
    if(::GetModuleFileName(NULL, path, sizeof(path)) == 0)
    {
	CORBA_String_var err = GetErrorText();
	cout << "GetModuleFileName failed: " << err << endl;
	return false;
    }
    
    const char* keyname;

    //
    // Add two keys EventMessageFile which contains the path to the
    // resources file, and TypesSupported which contains a DWORD with
    // the permitted event types.
    //
    keyname = "EventMessageFile";
    if(::RegSetValueEx(key, keyname, 0,
		       REG_EXPAND_SZ, (unsigned char*)path,
		       strlen(path)+1) == ERROR_SUCCESS)
    {
	DWORD supported = EVENTLOG_ERROR_TYPE|EVENTLOG_WARNING_TYPE|
	    EVENTLOG_INFORMATION_TYPE;
	keyname = "TypesSupported";
	if(::RegSetValueEx(key, "TypesSupported", 0, REG_DWORD,
			   (unsigned char*)&supported, sizeof(supported))
	   == ERROR_SUCCESS)
	{
	    return true;
	}
    }

    CORBA_String_var err = GetErrorText();
    cout << "RegSetValueEx " << keyname << " failed: " << err << endl;
    return false;
}

// ----------------------------------------------------------------------
// OBNTEventViewer constructor/destructor
// ----------------------------------------------------------------------

OBNTEventViewer::OBNTEventViewer(const char* service, DWORD eventId)
    : service_(CORBA_string_dup(service)),
      eventId_(eventId)
{
    eventSource_ = ::RegisterEventSource(NULL, service);
    if(eventSource_ == 0)
	throw CORBA_INITIALIZE("RegisterEventSource failed.");
}

OBNTEventViewer::~OBNTEventViewer()
{
    ::DeregisterEventSource(eventSource_);
}

// ----------------------------------------------------------------------
// OBNTEventViewer public member implementation
// ----------------------------------------------------------------------

OBNTEventViewer*
OBNTEventViewer::instance(const char* service, DWORD eventId)
{
    if(!instance_)
	instance_ = new OBNTEventViewer(service, eventId);
    return (OBNTEventViewer*)instance_;
}

bool
OBNTEventViewer::install()
{
    //
    // Constructor the key name
    //
    CORBA_String_var keyName =
	CORBA_string_dup("SYSTEM\\CurrentControlSet\\Services\\EventLog"
			 "\\Application\\");
    keyName += service_;

    //
    // The key goes into HKEY_LOCAL_MACHINE
    //
    HKEY regHandle = HKEY_LOCAL_MACHINE;

    //
    // First try to open the key.
    //
    HKEY keyHandle;
    if(::RegOpenKeyEx(regHandle, keyName, 0, KEY_ALL_ACCESS, &keyHandle)
       != ERROR_SUCCESS)
    {
	//
	// Create the key
	//
	DWORD d;
	if(::RegCreateKeyEx(regHandle, keyName, 0, "REG_SZ",
			    REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, 0,
			    &keyHandle, &d) != ERROR_SUCCESS)
	{
	    CORBA_String_var err = GetErrorText();
	    cout << "RegCreateKeyEx " << keyName << " failed: " << err << endl;
	    return false;
	}
    }

    bool rc = installValues(keyHandle);

    ::RegCloseKey(keyHandle);

    return rc;
}

bool
OBNTEventViewer::uninstall()
{
    //
    // Constructor the key name
    //
    CORBA_String_var keyName =
	CORBA_string_dup("SYSTEM\\CurrentControlSet\\Services\\EventLog"
			 "\\Application\\");
    keyName += service_;

    return ::RegDeleteKey(HKEY_LOCAL_MACHINE, keyName) == ERROR_SUCCESS;
}
    
void
OBNTEventViewer::message(const char* msg)
{
    emitMessage(EVENTLOG_INFORMATION_TYPE, msg);
}

void
OBNTEventViewer::error(const char* msg)
{
    emitMessage(EVENTLOG_ERROR_TYPE, msg);
}

void
OBNTEventViewer::warning(const char* msg)
{
    emitMessage(EVENTLOG_WARNING_TYPE, msg);
}

void
OBNTEventViewer::nca(const char* msg)
{
    emitMessage(EVENTLOG_ERROR_TYPE, msg);
}

//
// Define callbacks
//
static void WINAPI
_OB_serviceCtrl(DWORD ctrlCode)
{
    OBNTService::instance() -> control(ctrlCode);
}

static void WINAPI
_OB_serviceMain(DWORD argc, LPTSTR* argv)
{
    OBNTService::instance() -> main(argc, argv);
}

static BOOL WINAPI
_OB_controlHandler(DWORD dwCtrlType)
{
    switch( dwCtrlType )
    {
        case CTRL_BREAK_EVENT:  // use Ctrl+C or Ctrl+Break to simulate
        case CTRL_C_EVENT:      // SERVICE_CONTROL_STOP in debug mode
	{
	    //
	    // The service manager calls this in a seperate thread. So
	    // first we have inform JTC to adopt this thread.
	    //
	    JTCAdoptCurrentThread adopter;
            OBNTService::instance() -> stop();
            return TRUE;
            break;
	}
    }
    return FALSE;
}

//
// The OBNTService - only one instance is allowed
//
OBNTService* OBNTService::instance_ = 0;

// ----------------------------------------------------------------------
// OBNTService private and protected member implementation
// ----------------------------------------------------------------------

void
OBNTService::statusUpdate(DWORD currentState, DWORD exitCode, DWORD waitHint)
{
    //
    // when debugging we don't report to the SCM
    //
    if(debug_)
        return;
    
    if (currentState == SERVICE_START_PENDING)
        status_.dwControlsAccepted = 0;
    else
        status_.dwControlsAccepted = SERVICE_ACCEPT_STOP;

    status_.dwCurrentState = currentState;
    status_.dwWin32ExitCode = exitCode;
    status_.dwWaitHint = waitHint;

    if ( ( currentState == SERVICE_RUNNING ) ||
         ( currentState == SERVICE_STOPPED ) )
        status_.dwCheckPoint = 0;
    else
        status_.dwCheckPoint = checkPoint_++;

    //
    // Report the status of the service to the service control
    // manager.
    //
    if (!SetServiceStatus( statusHandle_, &status_))
        OBMessageViewer::instance() -> error("SetServiceStatus");
}

void
OBNTService::control(DWORD ctrlCode)
{
    //
    // Handle the requested control code.
    //
    switch(ctrlCode)
    {
    case SERVICE_CONTROL_STOP:
    {
        //
        // Stop the service.
        //
        // SERVICE_STOP_PENDING should be reported before setting the
        // Stop Event - hServerStopEvent - in ServiceStop().  This
        // avoids a race condition which may result in a 1053 - The
        // Service did not respond...  error.
        //
        statusUpdate(SERVICE_STOP_PENDING);
	//
	// The service manager calls this in a seperate thread. So
	// first we have inform JTC to adopt this thread.
	//
	JTCAdoptCurrentThread adopter;
        stop();
        return;
    }
        
    case SERVICE_CONTROL_INTERROGATE:
        //
        // Update the service status.
        //
        break;

    default:
        //
        // invalid control code
        //
        break;
    }

    statusUpdate(status_.dwCurrentState);
}

void
OBNTService::main(int argc, char** argv)
{
    //
    // Register our service control handler:
    //
    statusHandle_ = ::RegisterServiceCtrlHandler(name_, _OB_serviceCtrl);
    if(statusHandle_)
    {
        //
        // SERVICE_STATUS members that don't change in example
        //
        status_.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
        status_.dwServiceSpecificExitCode = 0;
        
        //
        // report the status to the service control manager.
        //
        statusUpdate(SERVICE_START_PENDING, NO_ERROR, 3000);
        start(argc, argv);
    }

    //
    // try to report the stopped status to the service control manager.
    //
    if (statusHandle_)
        statusUpdate(SERVICE_STOPPED);
}

// ----------------------------------------------------------------------
// OBNTService constructor and destructor
// ----------------------------------------------------------------------

OBNTService::OBNTService(const char* name, const char* title)
    : name_(CORBA_string_dup(name)),
      title_(CORBA_string_dup(title)),
      debug_(false),
      checkPoint_(0),
      useEventViewer_(false),
      eventId_(0)
{
    assert(instance_ == 0);
    instance_ = this;
}

OBNTService::OBNTService(const char* name, const char* title,
			 DWORD eventId)
    : name_(CORBA_string_dup(name)),
      title_(CORBA_string_dup(title)),
      debug_(false),
      checkPoint_(0),
      useEventViewer_(true),
      eventId_(eventId)
{
    assert(instance_ == 0);
    instance_ = this;
}

OBNTService::~OBNTService()
{
    instance_ = 0;
}

// ----------------------------------------------------------------------
// OBNTService public member implementation
// ----------------------------------------------------------------------

OBNTService*
OBNTService::instance()
{
    return instance_;
}

bool
OBNTService::install()
{
    bool rc = false;
    char path[512];

    //
    // Get path to executable.
    //
    if(::GetModuleFileName(NULL, path, sizeof(path)) == 0)
    {
        CORBA_String_var err = GetErrorText();
	cout << "GetModuleFileName failed: " << err << endl;
        return rc;
    }

    //
    // Open the service manager.
    //
    SC_HANDLE managerHandle =
	::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if(!managerHandle)
    {
        CORBA_String_var err = GetErrorText();
        cout << "OpenSCManager failed: " <<  err << endl;
        return rc;
    }

    //
    // Create a new service.
    //
    SC_HANDLE serviceHandle =
	::CreateService(managerHandle,              // SCManager database
			name_,                      // name of service
			title_,                     // name to display
			SERVICE_ALL_ACCESS,         // desired access
			SERVICE_WIN32_OWN_PROCESS,  // service type
			SERVICE_DEMAND_START,       // start type
			SERVICE_ERROR_NORMAL,       // error control type
			path,                       // service's binary
			NULL,                       // no load ordering group
			NULL,                       // no tag identifier
			"",                         // dependencies
			NULL,                       // LocalSystem account
			NULL);                      // no password
    
    if(serviceHandle)
    {
        cout << title_ << ": installed." << endl;
        ::CloseServiceHandle(serviceHandle);
        rc = true;
    }
    else
    {
        CORBA_String_var err = GetErrorText();
        cout << "CreateService Failed: " << err << endl;
    }
    
    ::CloseServiceHandle(managerHandle);

    if(rc && useEventViewer_)
    {
	//
	// Initialize the event viewer.
	//
	OBNTEventViewer* viewer = OBNTEventViewer::instance(name_, eventId_);
	rc = viewer -> install();
    }

    return rc;
}

bool
OBNTService::uninstall()
{
    bool rc = false;

    SC_HANDLE managerHandle =
        ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if(!managerHandle)
    {
        CORBA_String_var err = GetErrorText();
        cout << "OpenSCManager failed: " << err << endl;
        return rc;
    }

    SC_HANDLE serviceHandle = ::OpenService(managerHandle, name_,
					    SERVICE_ALL_ACCESS);
    if(serviceHandle)
    {
        //
        // try to stop the service
        //
        if(::ControlService(serviceHandle, SERVICE_CONTROL_STOP, &status_))
        {
            cout << "Stopping: " << title_ << endl;
            Sleep( 1000 );

            while(::QueryServiceStatus(serviceHandle, &status_))
            {
                if(status_.dwCurrentState == SERVICE_STOP_PENDING)
                {
                    cout << "." << flush;
                    Sleep( 1000 );
                }
                else
                    break;
            }
            if(status_.dwCurrentState == SERVICE_STOPPED)
                cout << title_ << ": stopped." << endl;
            else
                cout << title_ << ": failed to stop." << endl;
        }
        
        //
        // now remove the service
        //
        if(::DeleteService(serviceHandle))
        {
            rc = true;
            cout << title_ << ": removed." << endl;
        }
        else
        {
            CORBA_String_var err = GetErrorText();
            cout << "DeleteService failed: " << err << endl;
        }
        ::CloseServiceHandle(serviceHandle);
    }
    else
    {
        CORBA_String_var err = GetErrorText();
        cout << "OpenService failed: " << err << endl;
    }
    ::CloseServiceHandle(managerHandle);

    if(rc && useEventViewer_)
    {
	//
	// Initialize the event viewer.
	//
	OBNTEventViewer* viewer = OBNTEventViewer::instance(name_, eventId_);
	rc = viewer -> uninstall();
    }

    return rc;
}

void
OBNTService::run(int argc, char** argv)
{
    if(debug_)
    {
        ::SetConsoleCtrlHandler(_OB_controlHandler, TRUE);
        start(argc, argv);
    }
    else
    {
        //
        // Initialize the event viewer.
        //
	if(useEventViewer_)
	    OBNTEventViewer::instance(name_, eventId_);

        //
        // Define the service dispatch table.
        //
        SERVICE_TABLE_ENTRY dispatchTable[] =
        {
            { name_.inout(), (LPSERVICE_MAIN_FUNCTION)_OB_serviceMain },
            { NULL, NULL }
        };

        //
        // Start the service.
        //
        if(!::StartServiceCtrlDispatcher(dispatchTable))
	{
	    CORBA_String_var err =
		CORBA_string_dup("StartServiceCtrlDispatcher: ");
	    err += GetErrorText();
            OBMessageViewer::instance() -> error(err);
	}
    }
}
