73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
from vacuumDecay import *
|
|
from arc4 import ARC4
|
|
import copy
|
|
|
|
class KnownPlaintextAndKeylen(State, ABC):
|
|
def __init__(self, plaintext, ciphertext, keyLenBits, keyBits=None, turn=0, generation=0, playersNum=1, lastChange=None):
|
|
if keyBits==None:
|
|
keyBits = [0]*keyLenBits
|
|
self.turn = turn
|
|
self.generation = generation
|
|
self.keyBits = keyBits
|
|
self.keyLenBits = keyLenBits
|
|
self.plaintext = plaintext
|
|
self.ciphertext = ciphertext
|
|
self.lastChange = lastChange
|
|
self.decrypt = self._decrypt()
|
|
self.score = self.getScore()
|
|
|
|
def mutate(self, action):
|
|
newKeyBits = copy.copy(self.keyBits)
|
|
newKeyBits[action.data] = int(not newKeyBits[action.data])
|
|
return XorKnownPlaintextAndKeylen(self.plaintext, self.ciphertext, self.keyLenBits, newKeyBits, generation=self.generation+1, lastChange = action.data)
|
|
|
|
def getAvaibleActions(self):
|
|
for i in range(self.keyLenBits):
|
|
#if self.keyBits[i] == 0:
|
|
if self.lastChange != i:
|
|
yield Action(0, i)
|
|
|
|
def getKey(self):
|
|
s = ""
|
|
for i in range(int(self.keyLenBits/8)):
|
|
s += chr(int("".join([str(c) for c in self.keyBits[i*8:][:8]]),2))
|
|
return s
|
|
|
|
@abstractmethod
|
|
def _decrypt(self):
|
|
pass
|
|
|
|
def checkWin(self):
|
|
return self.decrypt == self.plaintext
|
|
|
|
def getScore(self):
|
|
diff = sum([bin(ord(a) ^ ord(b)).count("1") for a,b in zip(self.decrypt, self.plaintext)])
|
|
return diff / (len(self.plaintext)*8)
|
|
|
|
def __str__(self):
|
|
return "{"+self.getKey()+"}["+self.decrypt+"]"
|
|
|
|
def getTensor(self):
|
|
return torch.tensor(self.keyBits + list(map(int, ''.join([bin(ord(i)).lstrip('0b').rjust(8,'0') for i in self.decrypt]))))
|
|
|
|
def getModel(self):
|
|
pass
|
|
|
|
def getPriority(self, score):
|
|
return self.score + (1/self.keyLenBits)*0.01*self.generation
|
|
|
|
class XorKnownPlaintextAndKeylen(KnownPlaintextAndKeylen):
|
|
def _decrypt(self):
|
|
return ''.join(chr(ord(a) ^ ord(b)) for a,b in zip(self.ciphertext, self.getKey()))
|
|
|
|
class RC4KnownPlayintextAndKeylen(KnownPlaintextAndKeylen):
|
|
def _decrypt(self):
|
|
rc4 = ARC4(self.getKey())
|
|
return rc4.decrypt(self.ciphertext).decode("ascii")
|
|
|
|
if __name__=="__main__":
|
|
vd = WeakSolver(RC4KnownPlaintextAndKeylen())
|
|
|
|
# TODO:
|
|
# - Should use bytes for everything (not array of ints / string)
|