From 834bed6eb08e3ccabce2bf56f666c9f8dc61bd64 Mon Sep 17 00:00:00 2001
From: Ting-Jun Wang <levi900227@gmail.com>
Date: Fri, 22 Mar 2024 16:31:37 +0800
Subject: [PATCH] feat: debug mode & fix 'unix manager' asci in mail

- add debug mode. if in the debug mode, it will send the mails without
  any CD time. And it only use my Gmail account as the receiver so that
  the other admin will not receive this DEBUG mail.
- send the mail with HTML header so that we can use <pre> label to show
  the ascii emoji
---
 main.py | 55 ++++++++++++++++++++++++++++++-------------------------
 1 file changed, 30 insertions(+), 25 deletions(-)

diff --git a/main.py b/main.py
index c3b032f..ff8fc40 100644
--- a/main.py
+++ b/main.py
@@ -6,34 +6,37 @@ from bs4 import BeautifulSoup
 from email.mime.text import MIMEText
 import os
 
+DEBUG_MODE = True
+
 URL = 'https://www.cmlab.csie.ntu.edu.tw/status-gpu/'
 GPU_LIMIT = 2
 GOOGLE_CODE = os.environ['GOOGLE_CODE']
 MAIL_CD_HOUR = 6
 MAIL_MESSAGE = '''
-Hi, {}
-
-提醒您，您目前在 {} 伺服器上已經使用了 {} 張 {}，
-依照 CMLab 規定，每人在每台 CML workstation 上至多使用 2 顆 GPU，詳細規定請參閱 https://www.cmlab.csie.ntu.edu.tw/wiki/doku.php?id=workstation:rule
-
-為了公平起見，建議您降低 GPU 使用量！
-雖然我們不會直接處理，但若有人檢舉，我們會停止您運行的程式！
-若有其他問題歡迎來信討論！謝謝您的配合！
-
-* 信件由網管自動化工具寄出，若為錯誤寄送請忽略此信件。
-
-Best,
-CMLab Unix Manager, Ting-Jun Wang
-CMLab, National Taiwan University
-Email: unix_manager@cmlab.csie.ntu.edu.tw
-
-            __  __     _
-           / / / /__  (_)_ __
-   __  ___/ /_/ / _ \/ /\ \ /
+Hi, {}<br>
+<br>
+提醒您，您目前在 {} 伺服器上已經使用了 {} 張 {}，<br>
+依照 CMLab 規定，每人在每台 CML workstation 上至多使用 2 顆 GPU，詳細規定請參閱 https://www.cmlab.csie.ntu.edu.tw/wiki/doku.php?id=workstation:rule <br>
+<br>
+為了公平起見，建議您降低 GPU 使用量！<br>
+雖然我們不會直接處理，但若有人檢舉，我們會停止您運行的程式！<br>
+若有其他問題歡迎來信討論！謝謝您的配合！<br>
+<br>
+* 信件由網管自動化工具寄出，若為錯誤寄送請忽略此信件。<br>
+<br>
+Best,<br>
+CMLab Unix Manager, Ting-Jun Wang<br>
+CMLab, National Taiwan University<br>
+Email: unix_manager@cmlab.csie.ntu.edu.tw<br>
+<pre>
+            __  __     _             
+           / / / /__  (_)_ __        
+   __  ___/ /_/ / _ \/ /\ \ /        
   /  |/  /\____/_//_/_//_\_\____ ____
  / /|_/ / _ `/ _ \/ _ `/ _ `/ -_) __/
-/_/  /_/\_,_/_//_/\_,_/\_, /\__/_/
-                      /___/
+/_/  /_/\_,_/_//_/\_,_/\_, /\__/_/   
+                      /___/          
+</pre>
 '''
 
 def get_server_gpu_status() -> list:
@@ -131,6 +134,8 @@ def filter(server_status: dict, limit: int = 2) -> list:
 
 def mail_notify(server_name: str, gpu_name: str, violators: list) -> None:
     def check_send(log: dict, username: str) -> bool:
+        if DEBUG_MODE:
+            return True
         if username not in log:
             return True
         else:
@@ -155,11 +160,11 @@ def mail_notify(server_name: str, gpu_name: str, violators: list) -> None:
                 json.dump(send_log, fp)
             usage = violator['gpu']
             msg = MAIL_MESSAGE.format(username, server_name, usage, gpu_name)
-            msg = MIMEText(msg, 'plain') # 郵件內文
-            msg['Subject'] = f'[網管通知] 提醒您已經超過 CMLab GPU 使用限制！({username})'
+            msg = MIMEText(msg, 'html') # 郵件內文
+            msg['Subject'] = f'[網管通知] 提醒您已經超過 CMLab GPU 使用限制！({username})' + "(DEBUG MODE)" if DEBUG_MODE else ""
             msg['From'] = 'unix_manager@cmlab.csie.ntu.edu.tw'
-            msg['To'] = f'{username}@cmlab.csie.ntu.edu.tw'
-            msg['Cc'] = 'unix_manager@cmlab.csie.ntu.edu.tw'
+            msg['To'] = f'{username}@cmlab.csie.ntu.edu.tw' if not DEBUG_MODE else 'snsd0805@cmlab.csie.ntu.edu.tw'
+            msg['Cc'] = 'unix_manager@cmlab.csie.ntu.edu.tw' if not DEBUG_MODE else 'snsd0805@cmlab.csie.ntu.edu.tw'
 
             smtp = smtplib.SMTP('smtp.gmail.com', 587)
             smtp.ehlo()