Task: ADFGVX cipher




Go: package main

import (
    "bytes"
    "fmt"
    "io/ioutil"
    "log"
    "math/rand"
    "sort"
    "strings"
    "time"
)

var adfgvx = "ADFGVX"
var alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

func distinct(bs []byte) []byte {
    var u []byte
    for _, b := range bs {
        if !bytes.Contains(u, []byte{b}) {
            u = append(u, b)
        }
    }
    return u
}

func allAsciiAlphaNum(word []byte) bool {
    for _, b := range word {
        if !((b >= 48 && b <= 57) || (b >= 65 && b <= 90) || (b >= 97 && b <= 122)) {
            return false
        }
    }
    return true
}

func orderKey(key string) []int {
    temp := make([][2]byte, len(key))
    for i := 0; i < len(key); i++ {
        temp[i] = [2]byte{key[i], byte(i)}
    }
    sort.Slice(temp, func(i, j int) bool { return temp[i][0] < temp[j][0] })
    res := make([]int, len(key))
    for i := 0; i < len(key); i++ {
        res[i] = int(temp[i][1])
    }
    return res
}

func createPolybius() []string {
    temp := []byte(alphabet)
    rand.Shuffle(36, func(i, j int) {
        temp[i], temp[j] = temp[j], temp[i]
    })
    alphabet = string(temp)
    fmt.Println("6 x 6 Polybius square:\n")
    fmt.Println("  | A D F G V X")
    fmt.Println("---------------")
    p := make([]string, 6)
    for i := 0; i < 6; i++ {
        fmt.Printf("%c | ", adfgvx[i])
        p[i] = alphabet[6*i : 6*(i+1)]
        for _, c := range p[i] {
            fmt.Printf("%c ", c)
        }
        fmt.Println()
    }
    return p
}

func createKey(n int) string {
    if n < 7 || n > 12 {
        log.Fatal("Key should be within 7 and 12 letters long.")
    }
    bs, err := ioutil.ReadFile("unixdict.txt")
    if err != nil {
        log.Fatal("Error reading file")
    }
    words := bytes.Split(bs, []byte{'\n'})
    var candidates [][]byte
    for _, word := range words {
        if len(word) == n && len(distinct(word)) == n && allAsciiAlphaNum(word) {
            candidates = append(candidates, word)
        }
    }
    k := string(bytes.ToUpper(candidates[rand.Intn(len(candidates))]))
    fmt.Println("\nThe key is", k)
    return k
}

func encrypt(polybius []string, key, plainText string) string {
    temp := ""
outer:
    for _, ch := range []byte(plainText) {
        for r := 0; r <= 5; r++ {
            for c := 0; c <= 5; c++ {
                if polybius[r][c] == ch {
                    temp += fmt.Sprintf("%c%c", adfgvx[r], adfgvx[c])
                    continue outer
                }
            }
        }
    }
    colLen := len(temp) / len(key)
    
    if len(temp)%len(key) > 0 {
        colLen++
    }
    table := make([][]string, colLen)
    for i := 0; i < colLen; i++ {
        table[i] = make([]string, len(key))
    }
    for i := 0; i < len(temp); i++ {
        table[i/len(key)][i%len(key)] = string(temp[i])
    }
    order := orderKey(key)
    cols := make([][]string, len(key))
    for i := 0; i < len(key); i++ {
        cols[i] = make([]string, colLen)
        for j := 0; j < colLen; j++ {
            cols[i][j] = table[j][order[i]]
        }
    }
    res := make([]string, len(cols))
    for i := 0; i < len(cols); i++ {
        res[i] = strings.Join(cols[i], "")
    }
    return strings.Join(res, " ")
}

func decrypt(polybius []string, key, cipherText string) string {
    colStrs := strings.Split(cipherText, " ")
    
    maxColLen := 0
    for _, s := range colStrs {
        if len(s) > maxColLen {
            maxColLen = len(s)
        }
    }
    cols := make([][]string, len(colStrs))
    for i, s := range colStrs {
        var ls []string
        for _, c := range s {
            ls = append(ls, string(c))
        }
        if len(s) < maxColLen {
            cols[i] = make([]string, maxColLen)
            copy(cols[i], ls)
        } else {
            cols[i] = ls
        }
    }
    table := make([][]string, maxColLen)
    order := orderKey(key)
    for i := 0; i < maxColLen; i++ {
        table[i] = make([]string, len(key))
        for j := 0; j < len(key); j++ {
            table[i][order[j]] = cols[j][i]
        }
    }
    temp := ""
    for i := 0; i < len(table); i++ {
        temp += strings.Join(table[i], "")
    }
    plainText := ""
    for i := 0; i < len(temp); i += 2 {
        r := strings.IndexByte(adfgvx, temp[i])
        c := strings.IndexByte(adfgvx, temp[i+1])
        plainText = plainText + string(polybius[r][c])
    }
    return plainText
}

func main() {
    rand.Seed(time.Now().UnixNano())
    plainText := "ATTACKAT1200AM"
    polybius := createPolybius()
    key := createKey(9)
    fmt.Println("\nPlaintext :", plainText)
    cipherText := encrypt(polybius, key, plainText)
    fmt.Println("\nEncrypted :", cipherText)
    plainText2 := decrypt(polybius, key, cipherText)
    fmt.Println("\nDecrypted :", plainText2)
}




Python: 

from random import shuffle, choice
from itertools import product, accumulate
from numpy import floor, sqrt

class ADFGVX:
    
    def __init__(self, spoly, k, alph='ADFGVX'):
        self.polybius = list(spoly.upper())
        self.pdim = int(floor(sqrt(len(self.polybius))))
        self.key = list(k.upper())
        self.keylen = len(self.key)
        self.alphabet = list(alph)
        pairs = [p[0] + p[1] for p in product(self.alphabet, self.alphabet)]
        self.encode = dict(zip(self.polybius, pairs))
        self.decode = dict((v, k) for (k, v) in self.encode.items())

    def encrypt(self, msg):
        
        chars = list(''.join([self.encode[c] for c in msg.upper() if c in self.polybius]))
        colvecs = [(lett, chars[i:len(chars):self.keylen]) \
            for (i, lett) in enumerate(self.key)]
        colvecs.sort(key=lambda x: x[0])
        return ''.join([''.join(a[1]) for a in colvecs])

    def decrypt(self, cod):
        
        chars = [c for c in cod if c in self.alphabet]
        sortedkey = sorted(self.key)
        order = [self.key.index(ch) for ch in sortedkey]
        originalorder = [sortedkey.index(ch) for ch in self.key]
        base, extra = divmod(len(chars), self.keylen)
        strides = [base + (1 if extra > i else 0) for i in order]    
        starts = list(accumulate(strides[:-1], lambda x, y: x + y))  
        starts = [0] + starts                                        
        ends = [starts[i] + strides[i] for i in range(self.keylen)]  
        cols = [chars[starts[i]:ends[i]] for i in originalorder]     
        pairs = []                                                   
        for i in range((len(chars) - 1) // self.keylen + 1):
            for j in range(self.keylen):
                if i * self.keylen + j < len(chars):
                    pairs.append(cols[j][i])

        return ''.join([self.decode[pairs[i] + pairs[i + 1]] for i in range(0, len(pairs), 2)])


if __name__ == '__main__':
    PCHARS = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789')
    shuffle(PCHARS)
    POLYBIUS = ''.join(PCHARS)
    with open('unixdict.txt') as fh:
        WORDS = [w for w in (fh.read()).split() \
            if len(w) == 9 and len(w) == len(set(list(w)))]
        KEY = choice(WORDS)

    SECRET, MESSAGE = ADFGVX(POLYBIUS, KEY), 'ATTACKAT1200AM'
    print(f'Polybius: {POLYBIUS}, key: {KEY}')
    print('Message: ', MESSAGE)
    ENCODED = SECRET.encrypt(MESSAGE)
    DECODED = SECRET.decrypt(ENCODED)
    print('Encoded: ', ENCODED)
    print('Decoded: ', DECODED)