From 7556fe8988257bb6cdc029dd154d275d2c07f630 Mon Sep 17 00:00:00 2001
From: TLSM <duolsm@outlook.com>
Date: Mon, 6 Jun 2022 00:07:38 -0400
Subject: [PATCH] Add lottershe badge logic, badge helper.

---
 files/helpers/actions.py | 24 ++++++++++++++++++++++++
 files/helpers/lottery.py | 10 +++++++++-
 files/routes/admin.py    |  2 +-
 3 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 files/helpers/actions.py

diff --git a/files/helpers/actions.py b/files/helpers/actions.py
new file mode 100644
index 000000000..25fc50d8e
--- /dev/null
+++ b/files/helpers/actions.py
@@ -0,0 +1,24 @@
+from flask import g, abort
+from files.classes.user import User
+from files.classes.badges import Badge, BadgeDef
+
+# TODO: More sanity checks on passed parameters.
+# TODO: Add `replace=False` parameter which, when set true, removes any 
+#       existing badge with identical id & user and replaces with new one.
+def badge_grant(user_id, badge_id, desc='', url=''):
+	user = g.db.query(User).filter(User.id == int(user_id)).one_or_none()
+	if not user:
+		return False
+	elif user.has_badge(badge_id):
+		return True
+
+	badge = Badge(
+		badge_id=int(badge_id), 
+		user_id=user.id,
+		description=desc if desc != '' else None,
+		url=url if url != '' else None,
+	)
+
+	g.db.add(badge)
+	g.db.commit()
+	return True
diff --git a/files/helpers/lottery.py b/files/helpers/lottery.py
index 9ce9138ce..a523931b3 100644
--- a/files/helpers/lottery.py
+++ b/files/helpers/lottery.py
@@ -3,9 +3,11 @@ from random import choice
 from sqlalchemy import *
 from files.helpers.alerts import *
 from files.helpers.wrappers import *
+from files.helpers.actions import badge_grant
 from flask import g
 from .const import *
 
+LOTTERY_WINNER_BADGE_ID = 137
 
 def get_active_lottery():
     return g.db.query(Lottery).order_by(Lottery.id.desc()).filter(Lottery.is_active).one_or_none()
@@ -41,10 +43,16 @@ def end_lottery_session():
     winning_user = next(filter(lambda x: x.id == winner, participating_users))
     winning_user.coins += active_lottery.prize
     winning_user.total_lottery_winnings += active_lottery.prize
+    badge_grant(winner, LOTTERY_WINNER_BADGE_ID)
 
     for user in participating_users:
         chance_to_win = user.currently_held_lottery_tickets / len(raffle) * 100
-        notification_text = f'You won {active_lottery.prize} dramacoins in the lottery! Congratulations!\nOdds of winning: {chance_to_win}%' if user.id == winner else "You did not win the lottery. Better luck next time!\nOdds of winning: {chance_to_win}%"
+        if user.id == winner:
+            notification_text = f'You won {active_lottery.prize} dramacoins in the lottery! ' \
+                + f'Congratulations!\nOdds of winning: {chance_to_win}%' 
+        else:
+            notification_text = f'You did not win the lottery. Better luck next time!\n' \
+                + f'Odds of winning: {chance_to_win}%'
         send_repeatable_notification(user.id, notification_text)
         user.currently_held_lottery_tickets = 0
 
diff --git a/files/routes/admin.py b/files/routes/admin.py
index 5deed1fe2..15d4fc2f7 100644
--- a/files/routes/admin.py
+++ b/files/routes/admin.py
@@ -622,7 +622,7 @@ def badge_grant_post(v):
 	try: badge_id = int(request.values.get("badge_id"))
 	except: abort(400)
 
-	if badge_id in {16,17,94,95,96,97,98,109} and v.id != AEVANN_ID:
+	if badge_id in {16,17,94,95,96,97,98,109,137} and v.id != AEVANN_ID:
 		abort(403)
 
 	if user.has_badge(badge_id):