From f04c7540377bc392d2d08a096ff6f5330dc42687 Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Sat, 10 Nov 2018 21:16:39 +0000
Subject: [PATCH] Add AllowedDOmains property

---
 RhinoReminds/ClientListener.cs | 26 ++++++++++++++++++--------
 RhinoReminds/Program.cs        | 10 ++++++++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/RhinoReminds/ClientListener.cs b/RhinoReminds/ClientListener.cs
index e0d9d1a..9d40093 100644
--- a/RhinoReminds/ClientListener.cs
+++ b/RhinoReminds/ClientListener.cs
@@ -1,4 +1,5 @@
 ﻿using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Net;
@@ -22,6 +23,8 @@ namespace RhinoReminds
 		public string Hostname => Jid.Split('@')[1];
 		private readonly string password;
 
+		public readonly List<string> AllowedDomains = new List<string>();
+
 		public string ReminderFilePath { get; set; } = "./reminders.xml";
 		private ReminderList reminderList = new ReminderList();
 
@@ -62,6 +65,10 @@ namespace RhinoReminds
 		#region XMPP Event Handling
 		private bool subscriptionRequestHandler(Jid from)
 		{
+			if (!AllowedDomains.Contains("*") && !AllowedDomains.Contains(from.Domain)) {
+				sendChatMessage(from, "Sorry! The domain of your JID doesn't match the ones in my allowed list.");
+				return false;
+			}
 			Console.WriteLine($"[Rhino/SubscriptionRequest] Approving subscription from {from}");
 			return true;
 		}
@@ -90,6 +97,11 @@ namespace RhinoReminds
 
 		private void messageHandler(Message message)
 		{
+			if (!AllowedDomains.Contains("*") && !AllowedDomains.Contains(message.From.Domain)) {
+				sendChatMessage(message.From, "Sorry! The domain of your JID doesn't match the ones in my allowed list.");
+				return;
+			}
+
 			string messageText = message.Body;
 			string[] parts = Regex.Split(messageText.Trim(), @"\s+");
 			string instruction = parts[0].ToLower();
@@ -117,8 +129,6 @@ namespace RhinoReminds
 					break;
 
 				case "remind":
-					Console.WriteLine("[Rhino/Reciever] Identified remind request");
-
 					DateTime dateTime; string rawDateTimeString;
 					try {
 						dateTime = AIRecogniser.RecogniseDateTime(messageText, out rawDateTimeString);
@@ -183,10 +193,10 @@ namespace RhinoReminds
 			Reminder nextReminder = reminderList.GetNextReminder();
 			reminderList.OnReminderListUpdate += (object sender, Reminder newReminder) => {
 				Reminder newNextReminder = reminderList.GetNextReminder();
-				Console.WriteLine("[Rhino/Reminderd/Canceller] Reminder added - comparing.");
-				Console.WriteLine($"[Rhino/Reminderd/Canceller] {nextReminder} / {newNextReminder}");
+				//Console.WriteLine("[Rhino/Reminderd/Canceller] Reminder added - comparing.");
+				//Console.WriteLine($"[Rhino/Reminderd/Canceller] {nextReminder} / {newNextReminder}");
 				if (nextReminder != newNextReminder) {
-					Console.WriteLine($"[Rhino/Reminderd/Canceller] Cancelling");
+					//Console.WriteLine($"[Rhino/Reminderd/Canceller] Cancelling");
 					nextReminder = newNextReminder;
 					cancellationSource.Cancel();
 				}
@@ -199,15 +209,15 @@ namespace RhinoReminds
 					if (nextReminder != null) {
 						nextWaitingTime = nextReminder.Time - DateTime.Now;
 						if (DateTime.Now < nextReminder.Time) {
-							Console.WriteLine($"[Rhino/Reminderd] Sleeping for {nextWaitingTime}");
+							//Console.WriteLine($"[Rhino/Reminderd] Sleeping for {nextWaitingTime}");
 							await Task.Delay(nextWaitingTime, cancellationToken);
 						}
 					} else {
-						Console.WriteLine("[Rhino/Reminderd] Sleeping until interrupted");
+						//Console.WriteLine("[Rhino/Reminderd] Sleeping until interrupted");
 						await Task.Delay(Timeout.Infinite, cancellationToken);
 					}
 				} catch (TaskCanceledException) {
-					Console.WriteLine("[Rhino/Reminderd] Sleep interrupted, recalculating");
+					//Console.WriteLine("[Rhino/Reminderd] Sleep interrupted, recalculating");
 					cancellationSource = new CancellationTokenSource();
 					cancellationToken = cancellationSource.Token;
 					continue;
diff --git a/RhinoReminds/Program.cs b/RhinoReminds/Program.cs
index 2c079a1..f4253d0 100644
--- a/RhinoReminds/Program.cs
+++ b/RhinoReminds/Program.cs
@@ -9,6 +9,7 @@ namespace RhinoReminds
 {
 	public class ProgramSettings
 	{
+		public string AllowedDomain = "*";
 		public string Filepath = "./reminders.xml";
 		public string Jid = null;
 
@@ -30,6 +31,7 @@ namespace RhinoReminds
 				}
 
 				switch (args[i]) {
+					case "-h":
 					case "--help":
 						Console.WriteLine("--- RhinoReminds ---");
 						Console.WriteLine("> An XMPP reminder bot");
@@ -41,6 +43,7 @@ namespace RhinoReminds
 						Console.WriteLine("Options:");
 						Console.WriteLine("    -h  --help    Show this message");
 						Console.WriteLine($"    -f  --file    Specify where to save reminders (default: {settings.Filepath})");
+						Console.WriteLine("        --domain {domain}  Set the domain users are allowed to originate at. Defaults to any domain.");
 						Console.WriteLine();
 						Console.WriteLine("Environment Variables:");
 						Console.WriteLine("    XMPP_JID       The JID to login to");
@@ -50,6 +53,12 @@ namespace RhinoReminds
 					case "--jid":
 						settings.Jid = args[++i];
 						break;
+
+					case "-d":
+					case "--domain":
+						settings.AllowedDomain = args[++i];
+						break;
+
 				}
 			}
 
@@ -66,6 +75,7 @@ namespace RhinoReminds
 			ClientListener client = new ClientListener(settings.Jid, settings.Password) {
 				ReminderFilePath = settings.Filepath
 			};
+			client.AllowedDomains.Add(settings.AllowedDomain);
 			client.Start().Wait();
 		}
 	}
