.BANCHO. Set user flag at first login

databaseHelper.py

@@ -2,6 +2,7 @@ import pymysql
 import bcolors
 import consoleHelper
 import threading
+import glob
 
 class db:
 	"""A MySQL database connection"""

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)
 

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))

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])