diff --git a/AATC_Client.py b/AATC_Client.py index 6918b39..116a1f0 100644 --- a/AATC_Client.py +++ b/AATC_Client.py @@ -40,7 +40,7 @@ def Connect(remote_ip,PORT): -import socket,codecs,ast,recvall +import socket,codecs,ast,recvall,AATC_Crypto #Create Socket #Create Connection @@ -56,6 +56,7 @@ def Connect(remote_ip,PORT): class UserInterface: def __init__(self,Connection): self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con) self.Username = "" print("Welcome to the AATC connection interface") @@ -233,12 +234,12 @@ def Exit(self): ############################################## ############################################## def Send(self,Code,data): - Info = codecs.encode(str((Code,data))) + Info = self.Crypto.Encrypt(codecs.encode(str((Code,data)))) self.con.sendall(Info) def Recv(self): #Returns tuple of Sucess,Message,Data of which data may just be useless for that function try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # Sucess, Message , Data return data[0],data[1],data[2] diff --git a/AATC_Create_Graph.py b/AATC_Create_Graph.py index 6034a63..351d513 100644 --- a/AATC_Create_Graph.py +++ b/AATC_Create_Graph.py @@ -37,7 +37,7 @@ def drange(start,end,step): yRange = yEnd - yStart zRange = zEnd - zStart graph.Add_Edges(xRange,yRange,zRange) -graph.clean_edges() +##graph.clean_edges() #No longer nessesary graph.Build_Node_Cache() graph.SaveGraph() diff --git a/AATC_Crypto.py b/AATC_Crypto.py new file mode 100644 index 0000000..d79adb5 --- /dev/null +++ b/AATC_Crypto.py @@ -0,0 +1,115 @@ +#AATC crypto module +# Diffie-Hellman from 'pip install diffiehellman ' +import diffiehellman.diffiehellman as DH +import codecs,recvall,ast +from Crypto.Cipher import AES + +class Crypter: + """ + A class to encrypt and decrypt a connection with AES. Takes a connection object and mode and sets up a secret key for both participants. + This key is then used to create an AES object which encrypts and decrypts messages via the Encrypt and Decrypt functions. + These strings must be in binary format. The length of the string will be padded to a length multiple of 16 in size. + The object will communicate with another Crypter object via the socket and pass public keys. + Best for use with communication with string versions of python objects and resulting converstion using 'ast' as this will remove padding whitespace. + + If in 'CLIENT' mode the Crypter will start communication. It will send a message ('PublicKey,(Client_Public_Key,)) and the server will return the servers public key. + The CLIENT will then send an 'Exit' message to end communication. + + Per default the Crypter will assume it should generate the keys when __init__ is called. This can be disabled by creating the object with AutoGenerate set to False. + + Mode can be changed via SetMode. Normal options are 'CLIENT' and 'SERVER'. + + """ + def __init__(self, con, mode = "CLIENT",AutoGenerate = True): + self.con = con + self.SetMode(mode) + if AutoGenerate: + self.GenerateKey() + + def SetMode(self,mode): + self.mode = mode + + def GenerateKey(self,key_size = 256): + print("Generating encryption keys. Please stand by...") #Generating keys takes a long time and have found no way to shorted key length + if self.mode == "SERVER": + self.ServerGenerateKey() + elif self.mode == "CLIENT": + self.ClientGenerateKey(key_size) + else: + raise ValueError("Crypter: Incorrect mode set") + print("Encryption keys generated",self.shared_key) + + + def ServerGenerateKey(self): + DifHel = DH.DiffieHellman() + + + Exit = False + while not Exit : + data = self.Recv() + Command, Arguments = data[0],data[1] + + if Command == "PublicKey": + try: + DifHel.generate_public_key() + DifHel.generate_shared_secret(Arguments[0]) + self.Send(("PublicKey",(DifHel.public_key,))) + except: + self.Send((False,())) + + elif Command == "Exit": + self.Send((True,())) + Exit = True + + self.shared_key = DifHel.shared_key[0:32] + self.AES = AES.new(self.shared_key) + + + + + def ClientGenerateKey(self,key_size): + DifHel = DH.DiffieHellman() + DifHel.generate_public_key() + + + self.Send(("PublicKey",(DifHel.public_key,))) + data = self.Recv() + if data[0] == False: + raise Exception("Error occured while transmitting public key") + DifHel.generate_shared_secret(data[1][0]) + + self.Send(("Exit",())) + data = self.Recv() + if data[0] == False: + raise Exception("Server failed to commit to exit") + + self.shared_key = DifHel.shared_key[0:32] + self.AES = AES.new(self.shared_key) + + + + + def Encrypt(self,data): + pad_size = 16-len(data)%16 + data += b" "*pad_size + return self.AES.encrypt(data) + + def Decrypt(self,data): +## pad_size = 16-len(data)%16 # - should not be nessesary as data will come in 16 size blocks and change in message would thwart decryption. +## data += b" "*pad_size + return self.AES.decrypt(data) + + + + + def Send(self,data): + self.con.sendall(codecs.encode(str(data))) + def Recv(self): + try: + data = recvall.recvall(self.con) + data = ast.literal_eval(codecs.decode(data)) + # (Command,Arguments) + return data + #return data[0],data[1],data[2] + except Exception as e: + print(e) diff --git a/AATC_Drone.py b/AATC_Drone.py index 1fa51f9..77ab1d6 100644 --- a/AATC_Drone.py +++ b/AATC_Drone.py @@ -1,4 +1,4 @@ -import socket,codecs,ast,recvall,sys,heapq +import socket,codecs,ast,recvall,sys,heapq,AATC_Crypto from AATC_Coordinate import * @@ -25,15 +25,16 @@ class DroneInterface: """ def __init__(self,Connection): self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con) self.DroneName = "" print("Welcome to the AATC Connection interface") def Send(self,Code,data): - Info = codecs.encode(str((Code,data))) + Info = self.Crypto.Encrypt(codecs.encode(str((Code,data)))) self.con.sendall(Info) def Recv(self): #Returns tuple of Sucess,Message,Data of which data may just be useless for that function try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # Sucess, Message , Data return data[0],data[1],data[2] diff --git a/AATC_GPIO.py b/AATC_GPIO.py index a5cc03c..4407323 100644 --- a/AATC_GPIO.py +++ b/AATC_GPIO.py @@ -1,20 +1,24 @@ import threading,queue,time,random -import RPi.GPIO as GPIO +#import RPi.GPIO as GPIO -#GPIO.setmode(GPIO.BOARD) +##GPIO.setmode(GPIO.BOARD) -#GPIO.setup(11, GPIO.OUT) #red -#GPIO.setup(13, GPIO.OUT) #amber -#GPIO.setup(21, GPIO.OUT) #green -#GPIO.setup(26, GPIO.IN) #button +##GPIO.setup(11, GPIO.OUT) #red +##GPIO.setup(13, GPIO.OUT) #amber +##GPIO.setup(21, GPIO.OUT) #green +##GPIO.setup(26, GPIO.IN) #button class GPIO_Thread_Controller: def __init__(self,Command_Queue): + print("Creating Thread Controller") self.Command_Queue = Command_Queue self.Threads = {} self.Thread_Queues = {} def Create_Thread(self,Thread_Name): + if Thread_Name in Threads: #Close thread if already exists + self.Close_Thread(Thread_Name) + Thread_Queue = queue.Queue() thread = threading.Thread(target = GPIO_Thread,args = (Thread_Name,Thread_Queue)) self.Threads[Thread_Name] = thread @@ -151,7 +155,7 @@ def BlinkTest(Thread_Name,pin,frequency,cycles): #prints demonstration of blink time.sleep(pauseTime) return False -def PatternTest(Thread_Name, Pattern ,ReferenceTime=1): +def Pattern(Thread_Name, Pattern ,ReferenceTime=1): try: GPIO.setmode(GPIO.BOARD) GPIO.setup(11, GPIO.OUT) #red diff --git a/AATC_Monitor.py b/AATC_Monitor.py index b1d22c8..78333d4 100644 --- a/AATC_Monitor.py +++ b/AATC_Monitor.py @@ -2,7 +2,7 @@ # Used to display the flights of drones -import socket,codecs,ast,recvall,sys +import socket,codecs,ast,recvall,sys,AATC_Crypto class MonitorInterface: @@ -14,15 +14,16 @@ class MonitorInterface: """ def __init__(self,Connection): self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con) self.MonitorName = "" print("Welcome to the AATC Connection interface") def Send(self,Code,data): - Info = codecs.encode(str((Code,data))) + Info = self.Crypto.Encrypt(codecs.encode(str((Code,data)))) self.con.sendall(Info) def Recv(self): #Returns tuple of Sucess,Message,Data of which data may just be useless for that function try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # Sucess, Message , Data return data[0],data[1],data[2] diff --git a/AATC_NoFlyZoneGrapher.py b/AATC_NoFlyZoneGrapher.py index 9f7e1a4..c44dc6d 100644 --- a/AATC_NoFlyZoneGrapher.py +++ b/AATC_NoFlyZoneGrapher.py @@ -73,7 +73,7 @@ def Make_Values(self,NoFlyZoneData): for y in range(StartY,EndY+1): for z in range(StartZ,EndZ+1): NodeID = graph.Direct_NodeID(x,y,z) #Gets NodeID for that area - print(NodeID) + if NodeID in Values: v = max([Zone["Level"],Values[NodeID]]) else: diff --git a/AATC_Server_002.py b/AATC_Server_002.py index 6c800a9..2992747 100644 --- a/AATC_Server_002.py +++ b/AATC_Server_002.py @@ -1,4 +1,5 @@ -import codecs,ast,AATC_DB,socket,recvall,os,AATC_AStar,math,random,time,pickle +import codecs,ast,socket,recvall,os,math,random,time,pickle +import AATC_AStar,AATC_DB, AATC_Crypto from AATC_Coordinate import * def GetTime(): @@ -43,14 +44,15 @@ class UserConnection: def __init__(self,Connection): self.DB = AATC_DB.DBConnection() self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con, mode = "SERVER" ) self.UserID = -1 #Used to identify if has logged in yet self.NOFLYZONE_THRESHOLD_COST = 50 self.GRAPHFOLDERNAME = "GraphFolder" def Send(self,data): - self.con.sendall(codecs.encode(str(data))) + self.con.sendall(self.Crypto.Encrypt(codecs.encode(str(data)))) def Recv(self): try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # (Command,Arguments) return data @@ -470,12 +472,13 @@ class MonitorConnection: def __init__(self,Connection): self.DB = AATC_DB.DBConnection() self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con, mode = "SERVER") self.MonitorID = -1 #Used to identify if has logged in yet def Send(self,data): - self.con.sendall(codecs.encode(str(data))) + self.con.sendall(self.Crypto.Encrypt(codecs.encode(str(data)))) def Recv(self): try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # (Command,Arguments) return data @@ -662,12 +665,13 @@ class DroneConnection: def __init__(self,Connection): self.DB = AATC_DB.DBConnection() self.con = Connection + self.Crypto = AATC_Crypto.Crypter(self.con, mode = "SERVER") self.DroneID = -1 #Used to identify if has logged in yet def Send(self,data): - self.con.sendall(codecs.encode(str(data))) + self.con.sendall(self.Crypto.Encrypt(codecs.encode(str(data)))) def Recv(self): try: - data = recvall.recvall(self.con) + data = self.Crypto.Decrypt(recvall.recvall(self.con)) data = ast.literal_eval(codecs.decode(data)) # (Command,Arguments) return data