✨ Add option to set custom domain.

meashishsaini · meashishsaini · commit 9222835e636a · 2020-11-21T15:07:27.000+05:30
Signed-off-by: Ashish Saini &lt;sainiashish08@gmail.com&gt;
diff --git a/README.md b/README.md
@@ -11,11 +11,13 @@ Create and view temporary mailbox using 1secmail [API](https://www.1secmail.com/
 * _appdirs_
 
 ## Installation
-3. Install using `pip install git+https://github.com/meashishsaini/tmpmail-python`
+1. Install using `pip install git+https://github.com/meashishsaini/tmpmail-python`
 
 ## Usage
 ```
-tmpmail [-h] [-g [USERNAME]] [-r] [-t] [-b [BROWSER]] [id]
+tmpmail [-h] [-g [USERNAME]] [-r] [-t] [-b [BROWSER]]
+               [-d {1secmail.com,1secmail.org,1secmail.net,wwjmp.com,esiix.com}]
+               [id]
 
 positional arguments:
   id                    id of the email received
@@ -28,7 +30,9 @@ optional arguments:
   -t, --text            view email as pure text.
   -b [BROWSER], --browser [BROWSER]
                         open email in given browser.
+  -d {1secmail.com,1secmail.org,1secmail.net,wwjmp.com,esiix.com}, --domain {1secmail.com,1secmail.org,1secmail.net,wwjmp.com,esiix.com}
+                        set a custom domain supported by 1secmail.
 ```
 
 ## Credits
-The python version inspired by Siddharth Dushantha's [tmpmail](https://github.com/sdushantha/tmpmail) script.
+The python version is inspired by Siddharth Dushantha's [tmpmail](https://github.com/sdushantha/tmpmail) script.
diff --git a/tmpmail/api/one_sec_mail.py b/tmpmail/api/one_sec_mail.py
@@ -49,10 +49,20 @@ def dict_to_obj(cls, our_dict: dict):
 class OneSecMail(APIBase):
 	def __init__(self, username, domain="1secmail.org"):
 		super().__init__()
+		if not self.is_valid_domain(domain):
+			raise Exception("Illegal domain.")
 		self.domain = domain
 		self.username = username
 		self.url = f"https://www.1secmail.com/api/v1/?domain={self.domain}&login={self.username}"
 	
+	@classmethod
+	def is_valid_domain(cls, domain: str):
+		return domain in cls.valid_domains()
+
+	@classmethod
+	def valid_domains(cls):
+		return ["1secmail.com", "1secmail.org", "1secmail.net", "wwjmp.com", "esiix.com"]
+
 	def check_mailbox(self) -> list:
 		action = "getMessages"
 		check_url = f"{self.url}&action={action}"
diff --git a/tmpmail/main.py b/tmpmail/main.py
@@ -74,13 +74,21 @@ def parse():
 	argparser.add_argument("-r", "--recent", help="view most recent email.", action="store_true")
 	argparser.add_argument("-t", "--text", help="view email as pure text.", action="store_true")
 	argparser.add_argument("-b", "--browser", help="open email in given browser.", type=str, nargs="?", default=SUPPRESS)
+	argparser.add_argument("-d", "--domain", help="set a custom domain supported by 1secmail.", type=str, default=SUPPRESS, choices=OneSecMail.valid_domains())
 	args = argparser.parse_args()
 
 	config = load_config()
 
 	if "username" in args:
 		config["username"] = args.username if args.username else random_username()
 		create_config(username=config.get("username"), domain=config.get("domain"))
+	
+	if "domain" in args:
+		if OneSecMail.is_valid_domain(args.domain):
+			config["domain"] = args.domain
+			create_config(username=config.get("username"), domain=config.get("domain"))
+		else:
+			print("Illegal domain name provided.")
 
 	one_sec_mail = OneSecMail(**config)
 
