From 49b36e6acbe2803919f3e51e6ede434bdd25b34d Mon Sep 17 00:00:00 2001
From: Nyo <to@nyo.zz.mu>
Date: Sun, 1 May 2016 18:09:35 +0200
Subject: [PATCH] .BANCHO. Set user flag at first login

---
 databaseHelper.py | 1 +
 locationHelper.py | 8 ++++----
 loginEvent.py     | 7 ++++++-
 userHelper.py     | 9 +++++++++
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/databaseHelper.py b/databaseHelper.py
index 711bcc9..a4a584d 100644
--- a/databaseHelper.py
+++ b/databaseHelper.py
@@ -2,6 +2,7 @@ import pymysql
 import bcolors
 import consoleHelper
 import threading
+import glob
 
 class db:
 	"""A MySQL database connection"""
diff --git a/locationHelper.py b/locationHelper.py
index 1c09602..5038e37 100644
--- a/locationHelper.py
+++ b/locationHelper.py
@@ -5,7 +5,7 @@ import consoleHelper
 import bcolors
 
 # API URL
-url = "http://ip.zxq.co/"
+URL = "http://ip.zxq.co/"
 
 
 def getCountry(ip):
@@ -21,11 +21,11 @@ def getCountry(ip):
 
 	try:
 		# Try to get country from Pikolo Aul's Go-Sanic ip API
-		country = json.loads(urllib.request.urlopen("{}/{}".format(url, ip)).read().decode())["country"]
+		country = json.loads(urllib.request.urlopen("{}/{}".format(URL, ip)).read().decode())["country"]
 	except:
 		consoleHelper.printColored("[!] Error in get country", bcolors.RED)
 
-	return country
+	return country.upper()
 
 
 def getLocation(ip):
@@ -41,7 +41,7 @@ def getLocation(ip):
 
 	try:
 		# Try to get position from Pikolo Aul's Go-Sanic ip API
-		data = json.loads(urllib.request.urlopen("{}/{}".format(url, ip)).read().decode())["loc"].split(",")
+		data = json.loads(urllib.request.urlopen("{}/{}".format(URL, ip)).read().decode())["loc"].split(",")
 	except:
 		consoleHelper.printColored("[!] Error in get position", bcolors.RED)
 
diff --git a/loginEvent.py b/loginEvent.py
index 43ee8bb..101252b 100644
--- a/loginEvent.py
+++ b/loginEvent.py
@@ -126,7 +126,8 @@ def handle(flaskRequest):
 		if generalFunctions.stringToBool(glob.conf.config["server"]["localizeusers"]):
 			# Get location and country from IP
 			location = locationHelper.getLocation(requestIP)
-			country = countryHelper.getCountryID(locationHelper.getCountry(requestIP))
+			countryLetters = locationHelper.getCountry(requestIP)
+			country = countryHelper.getCountryID(countryLetters)
 		else:
 			# Set location to 0,0 and get country from db
 			print("[!] Location skipped")
@@ -137,6 +138,10 @@ def handle(flaskRequest):
 		responseToken.setLocation(location)
 		responseToken.setCountry(country)
 
+		# Set country in db if user has no country (first bancho login)
+		if userHelper.getCountry(userID) == "XX":
+			userHelper.setCountry(userID, countryLetters)
+
 		# Send to everyone our userpanel and userStats (so they now we have logged in)
 		glob.tokens.enqueueAll(serverPackets.userPanel(userID))
 		glob.tokens.enqueueAll(serverPackets.userStats(userID))
diff --git a/userHelper.py b/userHelper.py
index 9030135..6a4a756 100644
--- a/userHelper.py
+++ b/userHelper.py
@@ -275,3 +275,12 @@ def setAllowed(userID, allowed):
 	allowed -- allowed status. 1: normal, 0: banned
 	"""
 	glob.db.execute("UPDATE users SET allowed = ? WHERE id = ?", [allowed, userID])
+
+def setCountry(userID, country):
+	"""
+	Set userID's country (two letters)
+
+	userID -- userID
+	country -- country letters
+	"""
+	glob.db.execute("UPDATE users_stats SET country = ? WHERE id = ?", [country, userID])