simple_disconnect module for znc 1.6+

simple_disconnect.cpp

/*
 * Copyright (C) 2004-2015 ZNC, see the NOTICE file for details.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <znc/User.h>
#include <znc/IRCNetwork.h>

#define SIMPLE_DISCONNECT_DEFAULT_REASON "Auto disconnect on %s"
#define SIMPLE_DISCONNECT_DEFAULT_TIME   60


class CSimpleDisconnect;

class CSimpleDisconnectJob : public CTimer {
public:
	CSimpleDisconnectJob(CModule* pModule, unsigned int uInterval, unsigned int uCycles, const CString& sLabel, const CString& sDescription)
		: CTimer(pModule, uInterval, uCycles, sLabel, sDescription) {}

	virtual ~CSimpleDisconnectJob() {}

protected:
	void RunJob() override;
};

class CSimpleDisconnect : public CModule {
private:
	CString      m_sReason;
	unsigned int m_iDisconnectWait;
	bool         m_bClientSetDisconnect;
	bool         m_bWeSetDisconnect;

public:
	MODCONSTRUCTOR(CSimpleDisconnect) {
		m_sReason        = SIMPLE_DISCONNECT_DEFAULT_REASON;
		m_iDisconnectWait      = SIMPLE_DISCONNECT_DEFAULT_TIME;
		m_bClientSetDisconnect = false;
		m_bWeSetDisconnect     = false;

		AddHelpCommand();
		AddCommand("Reason", static_cast<CModCommand::ModCmdFunc>(&CSimpleDisconnect::OnReasonCommand), "[<text>]", "Prints or sets the disconnect reason (%s is replaced with the time you disconnected)");
		AddCommand("Timer", static_cast<CModCommand::ModCmdFunc>(&CSimpleDisconnect::OnTimerCommand), "", "Prints the current time to wait before disconnecting from the IRC server");
		AddCommand("SetTimer", static_cast<CModCommand::ModCmdFunc>(&CSimpleDisconnect::OnSetTimerCommand), "<seconds>", "Sets the time to wait before disconnecting from the IRC server");
		AddCommand("DisableTimer", static_cast<CModCommand::ModCmdFunc>(&CSimpleDisconnect::OnDisableTimerCommand), "", "Disables the wait time before disconnecting from the IRC server");
	}

	virtual ~CSimpleDisconnect() {}

	bool OnLoad(const CString& sArgs, CString& sMessage) override {
		CString sReasonArg;

		// Load DisconnectWait
		CString sFirstArg = sArgs.Token(0);
		if (sFirstArg.Equals("-notimer")) {
			SetDisconnectWait(0);
			sReasonArg = sArgs.Token(1, true);
		} else if (sFirstArg.Equals("-timer")) {
			SetDisconnectWait(sArgs.Token(1).ToUInt());
			sReasonArg = sArgs.Token(2, true);
		} else {
			CString sDisconnectWait = GetNV("disconnectwait");
			if (!sDisconnectWait.empty())
				SetDisconnectWait(sDisconnectWait.ToUInt(), false);
			sReasonArg = sArgs;
		}

		// Load Reason
		if (!sReasonArg.empty()) {
			SetReason(sReasonArg);
		} else {
			CString sSavedReason = GetNV("reason");
			if (!sSavedReason.empty())
				SetReason(sSavedReason, false);
		}

		// Set disconnect on load, required if loaded via webadmin
		if (GetNetwork()->IsIRCConnected() && !GetNetwork()->IsUserAttached())
			SetDisconnect(false);

		return true;
	}

	void OnIRCConnected() override {
		if (GetNetwork()->IsUserAttached())
			SetBack();
		else
			SetDisconnect(false);
	}

	void OnClientLogin() override {
		SetBack();
	}

	void OnClientDisconnect() override {
		/* There might still be other clients */
		if (!GetNetwork()->IsUserAttached())
			SetDisconnect();
	}

	void OnReasonCommand(const CString& sLine) {
		CString sReason = sLine.Token(1, true);

		if (!sReason.empty()) {
			SetReason(sReason);
			PutModule("Disconnect reason set");
		} else {
			PutModule("Disconnect reason: " + m_sReason);
			PutModule("Current disconnect reason would be: " + ExpandReason());
		}
	}

	void OnTimerCommand(const CString& sLine) {
		PutModule("Current timer setting: "
				+ CString(m_iDisconnectWait) + " seconds");
	}

	void OnSetTimerCommand(const CString& sLine) {
		SetDisconnectWait(sLine.Token(1).ToUInt());

		if (m_iDisconnectWait == 0)
			PutModule("Timer disabled");
		else
			PutModule("Timer set to "
					+ CString(m_iDisconnectWait) + " seconds");
	}

	void OnDisableTimerCommand(const CString& sLine) {
		SetDisconnectWait(0);
		PutModule("Timer disabled");
	}

	void SetDisconnect(bool bTimer = true) {
		if (bTimer) {
			RemTimer("simple_disconnect");
			AddTimer(new CSimpleDisconnectJob(this, m_iDisconnectWait, 1,
				"simple_disconnect", "Disconnects from the IRC server when all clients disconnect."));
		} else {
			if (!m_bClientSetDisconnect) {
				GetNetwork()->SetIRCConnectEnabled(false);
				PutIRC("QUIT :" + ExpandReason());
				m_bWeSetDisconnect = true;
			}
		}
	}

	void SetBack() {
		RemTimer("simple_disconnect");
		if (m_bWeSetDisconnect) {
			GetNetwork()->SetIRCConnectEnabled(true);
			PutIRC("CONNECT");
			m_bWeSetDisconnect = false;
		}
	}

private:
	CString ExpandReason() {
		CString sReason = m_sReason;
		if (sReason.empty())
			sReason = SIMPLE_DISCONNECT_DEFAULT_REASON;

		ExpandString(sReason);

		return sReason;
	}

/* Settings */
	void SetReason(CString& sReason, bool bSave = true) {
		if (bSave)
			SetNV("reason", sReason);
		m_sReason = sReason;
	}

	void SetDisconnectWait(unsigned int iDisconnectWait, bool bSave = true) {
		if (bSave)
			SetNV("disconnectwait", CString(iDisconnectWait));
		m_iDisconnectWait = iDisconnectWait;
	}

};

void CSimpleDisconnectJob::RunJob() {
	((CSimpleDisconnect*)GetModule())->SetDisconnect(false);
}

template<> void TModInfo<CSimpleDisconnect>(CModInfo& Info) {
	Info.SetWikiPage("simple_disconnect");
	Info.SetHasArgs(true);
	Info.SetArgsHelpText("You might enter up to 3 arguments, like -notimer disconnectmessage or -timer 5 disconnectmessage.");
}

NETWORKMODULEDEFS(CSimpleDisconnect, "This module will automatically disconnect from the IRC server when all clients disconnect from the bouncer.")

This module does not work for me unless I explicitly define flags (like -notimer Goodbye) - is this expected behaviour? Other than that, this module has saved me a lot of heartache :)