fixed l1q5

IS/Lab/Lab1/q5_john.py

@@ -1,20 +1,15 @@
-def main():
-    shift = (ord('C') - ord('y')) % 26
+from string import ascii_uppercase as A
 
-    ctext = "XVIEWYWI"
-    plaintext = []
+def idx(c): return A.index(c)
+def infer_shift(ct, pt):
+    return (idx(ct[0]) - idx(pt[0].upper())) % 26
 
-    for ch in ctext:
-        if ch.isalpha():
-            base = ord('a')
-            shifted = (ord(ch.lower()) - base - shift) % 26 + base
-            plaintext.append(chr(shifted))
-        else:
-            plaintext.append(ch)
+def decrypt(ct, shift):
+    return ''.join(A[(idx(c) - shift) % 26] if c.isalpha() else c for c in ct)
 
-    print(f"Attack type: Known plaintext attack")
-    print(f"Ciphertext: {ctext}")
-    print(f"Decrypted: {''.join(plaintext)}")
+known_ct = "CIW"
+known_pt = "yes"
+cipher = "XVIEWYWI"
 
-if __name__ == '__main__':
-    main()
+shift = infer_shift(known_ct, known_pt)
+print(decrypt(cipher, shift))