# Hill Climbing - M-209 - Known-Plaintext Attack
# HC_KPA v0.9 : 10/09/15 JFB (en Python2)
# HC_KPA v1.0 : 09/04/26 JFB (en Python3)
# based on George Lasry Algoritm - Cryptologia online - 

import string
import sys
import random
import math
import HC_SCORE
import m209
from HC_modif_KPC import *
from bib_util import *
import getopt
import time

theOverlaps = [
 [0, 0], [0, 1], [0, 2], [0, 3], [0, 4],
 [1, 0], [1, 5], [1, 6], [1, 7], [0, 8],
 [2, 1], [2, 5], [2, 9], [2,10], [2,11],
 [3, 2], [3, 6], [3, 9], [3,12], [3,13],
 [4, 3], [4, 7], [4,10], [4,12], [4,14],
 [5, 4], [5, 8], [5,11], [5,13], [5,14]
]

#======================================
def usage():
	print("usage:")
	print("  -c crypto	The cryptogram")
	print("  -M method	The method (BIG,TRI,...), default: ADE")
	print("  -m file_method	The file used with method")
	print("  -l letter	The letter (shift, default: Z")
	print("  -P pins	The Pins")
	print("  -L lugs	The lugs (ex: 1:5:4:13:2:7)")
	print("  -O overlaps	The overlaps (ex: 0:0:1...:1:1)")
	print("  -S list	Skip P(ins), S(wap), R(andom), L(ugs), O(verlaps)")
	print("  -t tries 	Number of tries, default: 15")
#======================================

#--------------------------------------------------
try:
  opts, args = getopt.getopt(sys.argv[1:], 
	"ho:c:M:m:t:l:L:S:P:O:v", ["help", "output=", "crypto=","method=",
	"method_file=", "tries=", "letter=", "lugs=", "skip=","pins=",
	"overlaps=", "verbose"])
except getopt.GetoptError as err:
  # print help information and exit:
  print(sys.argv[0],": ", str(err))# option -a not recognized"
  usage()
  sys.exit(2)

output = None
verbose = False
MAX_Z = 15
METHOD = "ADE"
file_method = "IC"
msg = ""
PINS = ""
LUGS = []
OVER = [0]*16
SKIP_LUGS = False
SKIP_PINS = False
SKIP_OVER = False
SKIP_SWAP = False
cle_letter_shift = "Z"

for o, a in opts:
  if o == "-v":
    verbose = True
  elif o in ("-h", "--help"):
    usage()
    sys.exit()
  elif o in ("-o", "--output"):
    output = a
  elif o in ("-c", "--crypto"):
    msg = litfic( a )
  elif o in ("-M", "--method"):
    METHOD = a
  elif o in ("-m", "--method_file"):
    file_method = a
  elif o in ("-t", "--tries"):
    MAX_Z = int(a)
  elif o in ("-l", "--letter"):
    cle_letter_shift = a
  elif o in ("-L", "--lugs"):
    LUGS=a.split(":")
    cle_lug=list( map(int,LUGS) )
  elif o in ("-S", "--skip"):
    les_algo = list(a)
    if 'L' in les_algo: SKIP_LUGS = True
    if 'P' in les_algo: SKIP_PINS = True
    if 'O' in les_algo: SKIP_OVER = True
    if 'W' in les_algo: SKIP_SWAP = True
    if 'R' in les_algo: SKIP_RAND = True
  elif o in ("-P", "--pins="):
    PINS = a
  elif o in ("-O", "--overlaps="):
    tb = a.split(":")
    OVER = list(map(int,tb))
    OVER.append(0)
  else:
    assert False, "unhandled option"
if msg == "" or file_method == "":
  usage()
  sys.exit(1)
HC_SCORE.score_init[METHOD]( file_method )
#--------------------------------------------------
e = m209.Hagelin ( cle_letter_shift, 0 )
MAX_SCORE = 0
nod = 0 # Number of Decipherment 
start_time = time.time()
for z in range(	MAX_Z ):
  # 1. Set the external wheels settings to AAAAAA
  e.rewind()

  # 2. Generate initial random pins and lugs settings
  #    (or fixes setting according to parameters)
  if LUGS:
    currentKeyLug = cle_lug
  else:
    currentKeyLug = initKeyLug()
  oldKeyLug = currentKeyLug

  if OVER:
    currentKeyOve = OVER
  else:
    currentKeyOve = [0] * 16
  oldKeyOve = currentKeyOve 

  if PINS:
    currentKeyPin = PINS
  else:
    currentKeyPin = initKeyPin()
  oldKeyPin = currentKeyPin

  #===> Configurer la cle
  e.setAllPinsAllW( currentKeyPin )
  e.setLugs( currentKeyLug )
  e.setOver( currentKeyOve )

  # --- 2.1 On dechiffre avec la cle courante
  plain = m209.chiffre( e, msg ); nod +=1

  # --- 2.2 on calcule le fitness
  HIGH_SCORE = HC_SCORE.score[METHOD]( plain )
  MAX_CLAIR = plain

  # 3. Repeat the following three loops (3a,3b,3c) until the ADE 
  #    no longer improves
  anyImprovement = True
  while anyImprovement:
    anyImprovement = False

    # 3a. Check all pins transformation as long as an improvement
    #     in ADE can be achieved
    # 3a1. 131 pins
    if not SKIP_PINS:
     improved = True
     while improved:
      improved = False
      e.setLugs( oldKeyLug )
      e.setOver( oldKeyOve )
      for i in range(131): 
        # --- External Key sets to "AAAAAA"
        e.rewind()

        # --- Toogle a Pin
        tb = list( oldKeyPin )
        if tb[i] == "1" : x = "0"
        else: x = "1"
        tb[ i ] = x
        currentKeyPin = "".join( tb )
        e.setAllPinsAllW( currentKeyPin )

        # --- To cipher
        plain = m209.chiffre( e, msg ) ; nod += 1

        # --- Calculate fitness
        score = HC_SCORE.score[METHOD]( plain )

        # --- Keep it if best
        if score > HIGH_SCORE :
          HIGH_SCORE = score			
          oldKeyPin = str(currentKeyPin)
          MAX_CLAIR = str(plain)
          anyImprovement = improved = True

    # 3a2. Hasard
    if not SKIP_RAND:
     improved = True
     while improved:
      improved = False
      e.setLugs( oldKeyLug )
      e.setOver( oldKeyOve )
      for i in range(131): 
        # --- External Key sets to "AAAAAA"
        e.rewind()

        # --- Toogle a Pin
        tb = list( oldKeyPin )
        if tb[i] == "1" : x = "0"
        else: x = "1"
        tb[ i ] = x
        currentKeyPin = "".join( tb )
        e.setAllPinsAllW( currentKeyPin )

        # --- To cipher
        plain = m209.chiffre( e, msg ) ; nod += 1

        # --- Calculate fitness
        score = HC_SCORE.score[METHOD]( plain )

        # --- Keep it if best
        if score > HIGH_SCORE :
          HIGH_SCORE = score			
          oldKeyPin = str(currentKeyPin)
          MAX_CLAIR = str(plain)
          anyImprovement = improved = True

     if not SKIP_SWAP:
       improved = True
       while improved:
         improved = False

     if not SKIP_SWAP:
       improved = True
       while improved:
         improved = False
         e.setLugs( oldKeyLug )
         e.setOver( oldKeyOve )
         k1 = 0 ; k2 = 1
         while k1 < 129 or k2 < 130:
           # --- External Key sets to "AAAAAA"
           e.rewind()

           # --- Swap a pair of pins
           tb = list( oldKeyPin )
           p1 = k1
           p2 = k2
           k2 = k2 + 1 
           if k2 == 130 : k1 = k1 + 1 ; k2 = k1 + 1
           if tb[p1] == tb[p2]: continue
           x = tb[p1]
           tb[p1] = tb[p2]
           tb[p2] = x
           currentKeyPin = "".join( tb )
           e.setAllPinsAllW( currentKeyPin )

           # --- To cipher
           plain = m209.chiffre( e, msg ) ; nod += 1

           # --- Calculate fitness
           score = HC_SCORE.score[METHOD]( plain )

           # --- Keep it if best
           if score > HIGH_SCORE :
             HIGH_SCORE = score			
             oldKeyPin = str(currentKeyPin)
             MAX_CLAIR = str(plain)
             anyImprovement = improved = True

    # 3b. Lugs (without overlap)
    if not SKIP_LUGS:
     e.setAllPinsAllW( oldKeyPin )
     e.setOver( oldKeyOve )
     improved = True
     while improved:
      improved = False
      for k in [0,1]:
        for i in range(5): 
          for j in range(i+1,6):
            # --- External Key sets to "AAAAAA"
            e.rewind()

            # --- Grow a Lug and dimish other
            tb = list( oldKeyLug )
            x = tb[i] ; y = tb[j]
            if k == 0 and x > 1 and y < 13:
              tb[i] -= 1 ; tb[j] += 1
            elif k == 1 and y > 1 and x < 13:
              tb[i] += 1 ; tb[j] -= 1
            else:
              continue
            currentKeyLug = list(tb)
            e.setLugs( currentKeyLug )

            # --- To cipher
            plain = m209.chiffre( e, msg ) ; nod += 1

            # --- Calculate fitness
            score = HC_SCORE.score[METHOD]( plain )

            # --- Keep it if best
            if score > HIGH_SCORE :
              HIGH_SCORE = score			
              oldKeyLug = list(currentKeyLug)
              MAX_CLAIR = str(plain)
              anyImprovement = improved = True

    # 3c. Lugs with overlap
    if not SKIP_OVER:
     e.setAllPinsAllW( oldKeyPin )
     e.setLugs( oldKeyLug )
     improved = True
     while improved:
      improved = False
      for k in [0,1]:
        for lug,ove in theOverlaps:
          # --- External Key sets to "AAAAAA"
          e.rewind()

          # --- Grow a Lug and grow a overlap
          t1 = list( oldKeyLug )
          t2 = list( oldKeyOve )
          if k == 0 and t1[lug] < 13  and t2[15] < 15:
            t1[lug] += 1 ; t2[ove] += 1 ; t2[15] += 1
          elif k == 1 and t1[lug] > 1 and t2[ove] > 0:
            t1[lug] -= 1 ; t2[ove] -= 1 ; t2[15] -= 1
          else:
            continue
          currentKeyLug = list(t1)
          currentKeyOve = list(t2)
          e.setLugs( currentKeyLug )
          e.setOver( currentKeyOve )

          # --- To cipher
          plain = m209.chiffre( e, msg ) ; nod += 1

          # --- Calculate fitness
          score = HC_SCORE.score[METHOD]( plain )

          # --- Keep it if best
          if score > HIGH_SCORE :
            HIGH_SCORE = score			
            oldKeyLug = list(currentKeyLug)
            oldKeyOve = list(currentKeyOve)
            MAX_CLAIR = str(plain)
            anyImprovement = improved = True
        

  if HIGH_SCORE > MAX_SCORE :
    elasp = time.time() - start_time
    MAX_SCORE = HIGH_SCORE
    print(z, ">>>MAX SCORE: ", HIGH_SCORE, "==", MAX_CLAIR[0:50])
    print("KEY: ", oldKeyPin, oldKeyLug, oldKeyOve, "NoD: ",  nod , elasp)
    sys.stdout.flush()
print("FIN---, Nod: ", nod, " elasp: ", time.time()-start_time)
