# Hill Climbing - M-209 - Cipher text Attack
# HC_SULLIVAN.py  v0.9 : 10/09/15 JFB 
# based on Geoff Sullivan Algorithm - Cryptologia (2002) - 

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
import copy

#======================================
def usage():
    print "usage:"
    print "  -c crypto	The cryptogram"
    print "  -M method	The method (BIG,TRI,...), default: IC"
    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 "  -S list	Skip P(ins), S(wap), L(ugs)"
    print "  -B         I choose the Best solution from first stage"
#======================================

WheelsPositions = [
    [  0, 26 ],
    [ 26, 25 ],
    [ 51, 23 ],
    [ 74, 21 ],
    [ 95, 19 ],
    [114, 17 ]
] 
TwoWheelsKick = [
  # r1,k1,r2,k2,   x, xl,   y, ylg
  [  0, 1, 1, 2 ],   # 0  wheels 1 (26) and 2 (25)
  [  0, 2, 1, 1 ],   # 1
  [  0, 1, 2, 2 ],   # 2  wheels 1 (26) and 3 (23)
  [  0, 2, 2, 1 ],   # 3
  [  0, 1, 3, 2 ],   # 4  wheels 1 (26) and 4 (21)
  [  0, 2, 3, 1 ],   # 5
  [  0, 1, 4, 2 ],   # 6  wheels 1 (26) and 5 (19)
  [  0, 2, 4, 1 ],   # 7
  [  0, 1, 5, 2 ],   # 8  wheels 1 (26) and 6 (17)
  [  0, 2, 5, 1 ],   # 9  
  [  1, 1, 2, 2 ],   #10  wheels 2 (25) and 3 (23)
  [  1, 2, 2, 1 ],   #11
  [  1, 1, 3, 2 ],   #12  wheels 2 (25) and 4 (21)
  [  1, 2, 3, 1 ],   #13
  [  1, 1, 4, 2 ],   #14  wheels 2 (25) and 5 (19)
  [  1, 2, 4, 1 ],   #15
  [  1, 1, 5, 2 ],   #16  wheels 2 (25) and 6 (17)
  [  1, 2, 5, 1 ],   #17
  [  2, 1, 3, 2 ],   #18  wheels 3 (23) and 4 (21)
  [  2, 2, 3, 1 ],   #19
  [  2, 1, 4, 2 ],   #20  wheels 3 (23) and 5 (19)
  [  2, 2, 4, 1 ],   #21
  [  2, 1, 5, 2 ],   #22  wheels 3 (23) and 6 (17)
  [  2, 2, 5, 1 ],   #23
  [  3, 1, 4, 2 ],   #24  wheels 4 (23) and 5 (21)
  [  3, 2, 4, 1 ],   #25
  [  3, 1, 5, 2 ],   #26  wheels 4 (23) and 6 (17)
  [  3, 2, 5, 1 ],   #27
  [  4, 1, 5, 2 ],   #28  wheels 5 (21) and 6 (17)
  [  4, 2, 5, 1 ]    #29  
]
#--------------------------------------------------
try:
  opts, args = getopt.getopt(sys.argv[1:], 
	"ho:c:M:m:l:L:S:P:vB", ["help", "output=", "crypto=","method=",
	"method_file=", "letter=", "lugs=", "skip=","pins=", "verbose", "best"])
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
best = False
MAX_Z = 15
METHOD = "IC"
file_method = "IC"
msg = ""
PINS = ""
LUGS = []
OVER = [0]*16
SKIP_LUGS = False
SKIP_PINS = False
SKIP_SWAP = False
cle_letter_shift = "Z"

for o, a in opts:
  if o == "-v":
    verbose = True
  elif o in ("-B", "--best"):
    best = 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 ("-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 'W' in les_algo: SKIP_SWAP = True
  elif o in ("-P", "--pins="):
    PINS = a
  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 )
nod = 0 # Number of Decipherment 
start_time = time.time()
#------------------ first Stage -------------------
print "============ first Stage ================="
solutions = [0]*30
MAX_Z = 0 ; MAX_Z_SCORE = 0
for z in range(	30 ):
  # 1. Set the external wheels settings to AAAAAA
  e.rewind()
  MAX_SCORE = 0

  # 2. Generate initial random pins 
  #    (or fixes setting according to parameters)
  (wh1, kick1, wh2, kick2) = TwoWheelsKick[z]
  (w1st, w1lg) = WheelsPositions[ wh1 ]
  (w2st, w2lg) = WheelsPositions[ wh2 ]

  LUGS=[0]*6
  LUGS[wh1]=kick1
  LUGS[wh2]=kick2
  currentKeyLug = LUGS
  oldKeyLug = currentKeyLug

  currentKeyOve = [0] * 16
  oldKeyOve = currentKeyOve

  currentKeyPin = initKeyPin()
  tb = list(currentKeyPin)
  for i in range(131):
    if not ((i>=w1st and i<(w1st+w1lg))or(i>=w2st and i<(w2st+w2lg))):
      tb[i] = "0"
  currentKeyPin = "".join( tb )
  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 steps until no longer improves
  anyImprovement = True
  while anyImprovement:
    anyImprovement = False

    # 3a. Check all pins transformation as long as an improvement
    if not SKIP_PINS:
     improved = True
     while improved:
      improved = False
      e.setLugs( oldKeyLug )
      e.setOver( oldKeyOve )

      # ---- 
      for iw1w2 in range(w1lg+w2lg): 
        # --- External Key sets to "AAAAAA"
        e.rewind()

        # --- Toogle a Pin
        if iw1w2 < w1lg :  i = w1st + iw1w2
        else:              i = w2st + iw1w2 - w1lg
        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
         e.setLugs( oldKeyLug )
         e.setOver( oldKeyOve )
         k1 = 0 ; k2 = 1
         kMax = w1lg + w2lg
         #while k1 < 129 or k2 < 130:
         while k1 < (kMax - 2) or k2 < (kMax - 1):
           # --- External Key sets to "AAAAAA"
           e.rewind()

           # --- Swap a pair of pins
           tb = list( oldKeyPin )
           if (k1 < w1lg):
             p1 = k1 + w1st
           else:
             p1 = k1 + w2st - w1lg
           if (k2 < w1lg):
             p2 = k2 + w1st
           else:
             p2 = k2 + w2st - w1lg
           k2 = k2 + 1 
           if k2 == (kMax - 1) : 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

  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 "CLE: ", oldKeyPin, oldKeyLug, oldKeyOve, "NoD: ",  nod , elasp
    sys.stdout.flush()
  if HIGH_SCORE > MAX_Z_SCORE:
    MAX_Z = z ; MAX_Z_SCORE = HIGH_SCORE
  copyOldKeyLug = copy.copy(oldKeyLug)
  copyOldKeyPin = copy.copy(oldKeyPin)
  solutions[z] = [ HIGH_SCORE, MAX_CLAIR[0:50], copyOldKeyPin, copyOldKeyLug ]

print "Best of Stage1: ", MAX_Z, ",", solutions[MAX_Z]
sys.stdout.flush()
print ">>>>>>>>>>>>>>> second stage <<<<<<<<<<<<<<<<<"
MAXMAX=0 
if best:
    MIN_Z = MAX_Z ; MAX_Z +=1
else:
    MIN_Z = 0 ; MAX_Z = 30
for z in range(MIN_Z,MAX_Z):
 print "z:", z, " ", solutions[z]
 #---------------- 
 WheelsToScan = [0,1,2,3,4,5]
 w1 = TwoWheelsKick[z][0]
 WheelsToScan.remove(w1)
 w2 = TwoWheelsKick[z][2]
 WheelsToScan.remove(w2)
 MAX_SCORE = 0
 MAX_CLAIR = ""
 while len(WheelsToScan) > 0:
  print "Wheels To Scan: ", WheelsToScan
  maxWheel = WheelsToScan[0]
  for w in range(len(WheelsToScan)):
   for nLugs in range(2,14):
    aW = WheelsToScan[w]
    oldKeyPin = copy.copy(solutions[z][2])
    oldKeyLug = copy.copy(solutions[z][3])
    oldKeyLug[aW] = nLugs      
    HIGH_SCORE= solutions[z][0]
    tb = list(oldKeyPin)
    wst = WheelsPositions[aW][0]
    wlg = WheelsPositions[aW][1]
    for i in range(wlg):
      if random.random() > 0.5: tb[i+wst]="1"
      else: tb[i+wst]="0"
    currentKeyPin = "".join( tb )
    oldKeyPin = currentKeyPin

    anyImprovement = True
    while anyImprovement:
      anyImprovement = False

      # 3a. Check all pins transformation as long as an improvement
      if not SKIP_PINS:
       improved = True
       while improved:
        improved = False
        e.setLugs( oldKeyLug )
        e.setOver( oldKeyOve )

        # ---- 
        for i in range(wst, wst+wlg):
          # --- 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
         e.setLugs( oldKeyLug )
         e.setOver( oldKeyOve )
         k1 = 0 ; k2 = 1
         kMax = wlg 
         while k1 < (kMax - 2) or k2 < (kMax - 1):
           # --- External Key sets to "AAAAAA"
           e.rewind()

           # --- Swap a pair of pins
           tb = list( oldKeyPin )
           p1 = k1 + wst
           p2 = k2 + wst
           k2 = k2 + 1
           if k2 == (kMax - 1) : 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

    if 1:
          elasp = time.time() - start_time
          print z, "===MAX_SCORE: ", HIGH_SCORE, "==", MAX_CLAIR[0:50],
          print "CLE: ", oldKeyPin, oldKeyLug, oldKeyOve, "NoD: ",  nod , elasp,
          print aW, WheelsToScan
          sys.stdout.flush()
          if HIGH_SCORE > MAX_SCORE :
            MAX_SCORE = HIGH_SCORE
            maxKeyLug = copy.copy( oldKeyLug )
            maxKeyPin = copy.copy( oldKeyPin )
            maxWheel  = aW
            MAXMAXCLR = copy.copy( MAX_CLAIR)
          if HIGH_SCORE > MAXMAX :
            MAXMAX = HIGH_SCORE
            maxmaxLug = copy.copy( oldKeyLug )
            maxmaxPin = copy.copy( oldKeyPin )
            maxmaxCLR = copy.copy( MAX_CLAIR )

  print "-------Wheel: ", maxWheel + 1, " Lugs: ", maxKeyLug[ maxWheel ], 
  print MAX_SCORE, maxKeyLug, maxKeyPin
  sys.stdout.flush()
  solutions[z][2] = copy.copy( maxKeyPin )
  solutions[z][3] = copy.copy( maxKeyLug )
  WheelsToScan.remove( maxWheel )
  
print "========>FIN<=========="
elasp = time.time() - start_time
print MAXMAX, maxmaxLug, maxmaxPin, maxmaxCLR[0:50], elasp


