import mysql.connector
from mysql.connector import errorcode
import json
from tkinter import *
import sys, traceback

'''
Demo to show that we can connect to database via python and follow on SQL injection attacks.  Creates a 
simple TKinter GUI for user input and output display.
Author: Tim Pierson, Dartmouth CS55, Winter 2021
	based on: Tim Pierson, Dartmouth CS61, Spring 2020


Requires installation of mysql connector: sudo pip install mysql-connector-python
MySQL connection based on: https://dev.mysql.com/doc/connector-python/en/connector-python-example-connecting.html
SQL injection based on https://www.youtube.com/watch?v=ciNHn38EyRc

Also required TKinter: sudo apt-get install python3-tk


Set global variable use_secure_call:
	False: use user input directly in SQL call, creating opportunity for SQL injection attack 
	True: use prepared statements to prevent SQL injection
'''


use_secure_call = False; #use prepared statement if true, else use user input directly (asking for trouble!)
 
	


def get_data():
	#Fetch database from database, make either a secure call or naive call based on use_secure_call variable
	#gets a cursor after making call, calls update_display to update GUI with results.
	#Input: None
	#Output: Updated GUI.

	#get database connection
	try:
		#make connection to database on host server using credentials provided
		cnx = mysql.connector.connect(user="root", password="seedubuntu", 
									host="localhost", database="cs55")
	#catch exceptions, alert user to problem, then exit
	except mysql.connector.Error as err:
		if err.errno == errorcode.ER_ACCESS_DENIED_ERROR:
			print("Something is wrong with your user name or password")
		elif err.errno == errorcode.ER_BAD_DB_ERROR:
			print("Database does not exist")
		else:
			print(err)
		cnx.close()
		exit()

	#get name of restaurant input by user in restaurant_name textbox
	restaurant = restaurant_name.get() 

	#make a secure or naive call to database depending on use_secure_call setting
	if not use_secure_call:
		cursor = fetch_naive(cnx,restaurant)
	else:
		cursor = fetch_secure(cnx,restaurant)

	#update display with data returned, then close connection
	update_display(cursor)
	cursor.close()
	cnx.close()


def fetch_naive(cnx,restaurant):
	#Query database for restaurant data using user input directly in SQL command.  This opens us up for SQL injection
	#Input: 
	#	cnx: connection to the database
	#	restaurant: is name of restaurant about which to fetch data
	#Output: cursor of data returned about restaurant.

	try:
		#use connection set up in main to fetch data from database in a cursor
		cursor = cnx.cursor()
		#include user input into SQL query -- this is asking for trouble!
		query = ("SELECT RestaurantName AS `Restaurant Name`, Address, Boro  " 
				+"FROM Restaurants " 
				+"WHERE RestaurantName LIKE '%" + restaurant +"%' LIMIT 20") #embedding user input into SQL command
		print("Query is:",query)
		cursor.execute(query)
		return cursor

	except: 
		print("An exception occurred in fetch_naive")
		#traceback.print_exc(file=sys.stdout)
		cnx.close()


def fetch_secure(cnx,restaurant):
	#Query database for restaurant data using prepared statement
	#Input: Name of restaurant about which to fetch data
	#Output: cursor of data returned about restaurant.

	try:
		#use connection set up in main to fetch data from database in a cursor
		cursor = cnx.cursor()
		#use prepared statement to avoid SQL injection attacks
		query = ("SELECT RestaurantName AS `Restaurant Name`, Address, Boro  " 
				+"FROM Restaurants "
				+"WHERE RestaurantName LIKE %s LIMIT 20") #note using %s instead of embedding user input
		cursor.execute(query, ('%'+restaurant+'%',)) #second param fills in for %s in command, must be python tuple
		return cursor

	except: 
		print("An exception occurred in fetch_secure")
		traceback.print_exc(file=sys.stdout)
		cnx.close()
		exit()


def update_display(cursor):
	#Update GUI with results from database query
	#Input: cursor of database query results
	#Output: updated GUI with database results

	#remove old labels at row 4 and greater
	for label in window.grid_slaves():
		if int(label.grid_info()["row"]) >= 4:
			label.grid_forget()
	#add new column headers with attribute names
	cols = cursor.description #column names plus other attributes stored in description, one tuple per column
	column_names = [] #store column names in this list
	for i,c in enumerate(cols): #loop over all columns
		Label(window,text=c[0]).grid(row=3,column=i,sticky=W)#
		column_names.append(c[0]) #column name is first item in tuple, extract it and append to column_names list
	print(column_names)

	#loop over data and create a label for each data row
	for i,row in enumerate(cursor):
		print(row)
		#loop over all columns and make entry in display grid
		for j in range(len(column_names)):
			Label(window,text=str(row[j]),bg="white").grid(row=4+i,column=j,sticky=W)



if __name__ == '__main__':
	#call get_data whenever search button is pressed
	window = Tk() 
	window.title('SQL Injection Attack Demo') 
	Label(window, text="Search restaurants").grid(row=0,pady=10)
	restaurant_name = Entry(window)
	restaurant_name.grid(row=0, column=1)
	search_button = Button(window, text="Search", command=get_data)
	search_button.grid(row=1,column=1)
	Label(window,text="No results").grid(row=4)
	window.mainloop()