From a97d7f80b0e9ad2ebdc8d8ca5f0539131a029514 Mon Sep 17 00:00:00 2001
From: girishsrivatsa <girishs@newzera.com>
Date: Thu, 4 Feb 2021 16:46:42 +0530
Subject: [PATCH] updated remove punctuation with key based character filtering

---
 pycipher/adfgvx.py     |  2 +-
 pycipher/autokey.py    |  4 ++--
 pycipher/base.py       | 25 +++++++++++++------------
 pycipher/delastelle.py |  4 ++--
 pycipher/foursquare.py | 14 +++++++-------
 pycipher/playfair.py   | 12 ++++++------
 pycipher/polybius.py   |  5 +++--
 7 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/pycipher/adfgvx.py b/pycipher/adfgvx.py
index 1c87185..0f8c0b7 100644
--- a/pycipher/adfgvx.py
+++ b/pycipher/adfgvx.py
@@ -53,4 +53,4 @@ def decipher(self,string):
         return step1    
 
 if __name__ == '__main__': 
-    print('use "import pycipher" to access functions')
+    print('use "import pycipher" to access functions')
\ No newline at end of file
diff --git a/pycipher/autokey.py b/pycipher/autokey.py
index 5bfdb50..492561b 100644
--- a/pycipher/autokey.py
+++ b/pycipher/autokey.py
@@ -31,8 +31,8 @@ def encipher(self,string):
         ret = ''
         for (i,c) in enumerate(string):
             if i<len(self.key): offset = self.a2i(self.key[i])
-            else: offset = self.a2i(string[i-len(self.key)])     
-            ret += self.i2a(self.a2i(c)+offset)
+            else: offset = self.a2i(string[i-len(self.key)])
+            ret += self.i2a(self.a2i(c)+offset,lower=c.islower())
         return ret    
 
     def decipher(self,string):
diff --git a/pycipher/base.py b/pycipher/base.py
index 2191285..ea2ece5 100644
--- a/pycipher/base.py
+++ b/pycipher/base.py
@@ -4,7 +4,8 @@
 Author: James Lyons
 Created: 2012-04-28
 '''
-import re
+import string
+import warnings
 
 class Cipher(object):
     def encipher(self,string):
@@ -14,16 +15,16 @@ def decipher(self,string):
         return string
         
     def a2i(self,ch):
-        ch = ch.upper()
-        arr = {'A':0,'B':1,'C':2,'D':3,'E':4,'F':5,'G':6,'H':7,'I':8,'J':9,'K':10,
-           'L':11,'M':12,'N':13,'O':14,'P':15,'Q':16,'R':17,'S':18,'T':19,'U':20,
-           'V':21,'W':22,'X':23,'Y':24,'Z':25}
-        return arr[ch]
+        return ord(ch)-65 if ch.isupper() else ord(ch)-97
 
-    def i2a(self,i):
-        i = i%26
-        arr = ('A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z')
-        return arr[i]
+    def i2a(self,i,lower=False):
+        i=i%26
+        return chr(i+97) if lower else chr(i+65)
         
-    def remove_punctuation(self,text,filter='[^A-Z]'):
-        return re.sub(filter,'',text.upper())
+    def remove_punctuation(self,text,characters=None):
+        no_punctuation = text.translate(str.maketrans("","",string.punctuation))
+        if characters:
+            filtered_chars = ''.join(ch for ch in no_punctuation if ch in characters)
+            if len(filtered_chars) != no_punctuation:
+                warnings.warn(message="The data contains characters missing in key which will be removed",category=ResourceWarning)
+        return no_punctuation
diff --git a/pycipher/delastelle.py b/pycipher/delastelle.py
index 3847355..abb7e4a 100644
--- a/pycipher/delastelle.py
+++ b/pycipher/delastelle.py
@@ -33,10 +33,10 @@ def encipher(self,string):
         :param string: The string to encipher.
         :returns: The enciphered string. The ciphertext will be 3 times the length of the plaintext.
         """           
-        string = self.remove_punctuation(string,filter='[^'+self.key+']')        
+        string = self.remove_punctuation(string,characters=self.key)
         ctext = ""
         for c in string:
-            ctext += ''.join([str(i) for i in L2IND[c]])
+            ctext += ''.join([str(i) for i in self.L2IND[c]])
         return ctext
 
     def decipher(self,string):
diff --git a/pycipher/foursquare.py b/pycipher/foursquare.py
index 8f9118e..00be6a0 100644
--- a/pycipher/foursquare.py
+++ b/pycipher/foursquare.py
@@ -22,13 +22,13 @@ def __init__(self,key1='zgptfoihmuwdrcnykeqaxvsbl',key2='mfnbdcrhsaxyogvituewlqz
         assert len(self.key2)==25, 'key2 is not length 25'
         
     def encipher_pair(self,a,b):
-        arow,acol = self.alph.index(a)/5, self.alph.index(a)%5
-        brow,bcol = self.alph.index(b)/5, self.alph.index(b)%5
+        arow,acol = self.alph.index(a)//5, self.alph.index(a)%5
+        brow,bcol = self.alph.index(b)//5, self.alph.index(b)%5
         return (self.key1[arow*5+bcol], self.key2[brow*5+acol])
         
     def decipher_pair(self,a,b):
-        arow,acol = self.key1.index(a)/5, self.key1.index(a)%5
-        brow,bcol = self.key2.index(b)/5, self.key2.index(b)%5
+        arow,acol = self.key1.index(a)//5, self.key1.index(a)%5
+        brow,bcol = self.key2.index(b)//5, self.key2.index(b)%5
         return (self.alph[arow*5+bcol], self.alph[brow*5+acol])
         
     def encipher(self,string):
@@ -42,7 +42,7 @@ def encipher(self,string):
         :param string: The string to encipher.
         :returns: The enciphered string.
         """                
-        string = self.remove_punctuation(string)  
+        string = self.remove_punctuation(string,characters=self.alph).upper()
         if len(string)%2 == 1: string = string + 'X'
         ret = ''
         for c in range(0,len(string.upper()),2):
@@ -61,10 +61,10 @@ def decipher(self,string):
         :param string: The string to decipher.
         :returns: The deciphered string.
         """              
-        string = self.remove_punctuation(string)  
+        string = self.remove_punctuation(string).upper()
         if len(string)%2 == 1: string = string + 'X'
         ret = ''
-        for c in range(0,len(string.upper()),2):
+        for c in range(0,len(string),2):
             a,b = self.decipher_pair(string[c],string[c+1])
             ret += a + b
         return ret    
diff --git a/pycipher/playfair.py b/pycipher/playfair.py
index 3689647..1130024 100644
--- a/pycipher/playfair.py
+++ b/pycipher/playfair.py
@@ -21,8 +21,8 @@ def __init__(self, key='ABCDEFGHIKLMNOPQRSTUVWXYZ'):
     def encipher_pair(self, a, b):
         if a == b:
             b = 'X'
-        arow, acol = int(self.key.index(a) / 5), self.key.index(a) % 5
-        brow, bcol = int(self.key.index(b) / 5), self.key.index(b) % 5
+        arow, acol = self.key.index(a) // 5, self.key.index(a) % 5
+        brow, bcol = self.key.index(b) // 5, self.key.index(b) % 5
         if arow == brow:
             return self.key[arow * 5 + (acol + 1) % 5] + self.key[brow * 5 + (bcol + 1) % 5]
         elif acol == bcol:
@@ -32,8 +32,8 @@ def encipher_pair(self, a, b):
         
     def decipher_pair(self, a, b):
         assert a != b, 'two of the same letters occurred together, illegal in playfair'
-        arow, acol = int(self.key.index(a) / 5), self.key.index(a) % 5
-        brow, bcol = int(self.key.index(b) / 5), self.key.index(b) % 5
+        arow, acol = self.key.index(a) // 5, self.key.index(a) % 5
+        brow, bcol = self.key.index(b) // 5, self.key.index(b) % 5
         if arow == brow:
             return self.key[arow * 5 + (acol - 1) % 5] + self.key[brow * 5 + (bcol - 1) % 5]
         elif acol == bcol:
@@ -52,7 +52,7 @@ def encipher(self, string):
         :param string: The string to encipher.
         :returns: The enciphered string.
         """    
-        string = self.remove_punctuation(string)  
+        string = self.remove_punctuation(string,characters=self.key).upper()
         string = re.sub(r'[J]', 'I', string)
         if len(string) % 2 == 1:
             string += 'X'
@@ -72,7 +72,7 @@ def decipher(self, string):
         :param string: The string to decipher.
         :returns: The deciphered string.
         """    
-        string = self.remove_punctuation(string)  
+        string = self.remove_punctuation(string).upper()
         if len(string) % 2 == 1:
             string += 'X'
         ret = ''
diff --git a/pycipher/polybius.py b/pycipher/polybius.py
index 2f5a7f5..4afb470 100644
--- a/pycipher/polybius.py
+++ b/pycipher/polybius.py
@@ -5,6 +5,7 @@
 '''
 from pycipher.base import Cipher
 import re
+import string
 
 ####################################################################################
 class PolybiusSquare(Cipher):
@@ -18,7 +19,7 @@ class PolybiusSquare(Cipher):
     """    
     def __init__(self,key='phqgiumeaylnofdxkrcvstzwb',size=5,chars=None):
         self.key = ''.join([k.upper() for k in key])
-        self.chars = chars or 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'[:size]
+        self.chars = chars or string.ascii_uppercase[:size]
         self.size = size
         assert len(self.key)==size*size, 'invalid key in init: must have length size*size, has length '+str(len(key))
         assert len(self.chars)==size, 'invalid chars in init: must have length=size, has length '+str(len(chars))
@@ -43,7 +44,7 @@ def encipher(self,string):
         :param string: The string to encipher.
         :returns: The enciphered string. The ciphertext will be twice the length of the plaintext.
         """           
-        string = self.remove_punctuation(string)#,filter='[^'+self.key+']')
+        string = self.remove_punctuation(string.upper(),self.key).upper()#,filter='[^'+self.key+']')
         ret = ''
         for c in range(0,len(string)):
             ret += self.encipher_char(string[c])