enc.py

import unittest
from collections import namedtuple


DIGRAM_TABLE = [
    ['E', 'S', 'P', 'I', 'A'],
    ['O', 'B', 'C', 'D', 'F'],
    ['G', 'H', 'K', 'L', 'M'],
    ['N', 'Q', 'R', 'T', 'U'],
    ['V', 'W', 'X', 'Y', 'Z']
]


Position = namedtuple('Position', "line column")


def make_digram(message, n, default_extra='W'):
    _msg = message.upper().replace(' ', '')
    digram = [_msg[i:i+n] for i in range(0, len(_msg), n)]
    for d in digram:
        if d.count(d[0]) > 1:
            el = d
            new_el = el[0] + default_extra + el[1]
            digram[digram.index(d)] = new_el
            msg = ''.join(digram)
            return make_digram(msg, n=2)
    return digram


def message_to_digram(message, n=2, default_extra='W'):
    if not message:
        raise ValueError('Message cant be empty...')

    last_position = -1
    digram = make_digram(message, n)
    if len(digram[last_position]) == 1:
        c = digram[last_position]
        digram[last_position] = c+default_extra

    return digram


def get_position(e):
    for l in DIGRAM_TABLE:
        for c in l:
            if e == c:
                return Position(DIGRAM_TABLE.index(l), l.index(c))
    return None


def find_positions(par):

    if len(par) != 2:
        raise ValueError("Argument error")

    return [get_position(e) for e in par]


def same_line(positions):
    lines = [p.line for p in positions]
    return len(set(lines)) == 1


def same_column(positions):
    columns = [p.column for p in positions]
    return len(set(columns)) == 1


def cypher_same_line(positions):
    line_len_diagram = len(DIGRAM_TABLE[0]) - 1
    cyphers = list()
    initial_pos = 0
    for p in positions:
        if p.column == line_len_diagram:
            cyphers.append(DIGRAM_TABLE[p.line][initial_pos])
        else:
            cyphers.append(DIGRAM_TABLE[p.line][p.column+1])

    return ''.join(cyphers)


def cypher_same_column(positions):
    column_len_diagram = len(DIGRAM_TABLE) - 1
    cyphers = list()
    initial_pos = 0
    for p in positions:
        if p.line == column_len_diagram:
            cyphers.append(DIGRAM_TABLE[initial_pos][p.column])
        else:
            cyphers.append(DIGRAM_TABLE[p.line + 1][p.column])

    return ''.join(cyphers)


def cypher_diff_line_column(positions):
    cyphers = list()
    p1, p2 = (p for p in positions)

    cyphers.append(DIGRAM_TABLE[p1.line][p2.column])
    cyphers.append(DIGRAM_TABLE[p2.line][p1.column])

    return ''.join(cyphers)


def cypher(positions):
    if same_line(positions):
        return cypher_same_line(positions)
    elif same_column(positions):
        return cypher_same_column(positions)
    else:
        return cypher_diff_line_column(positions)


def encrypt_playfair(str_in, act):
    if act == 'ENC':
        msg = message_to_digram(str_in)
        positions = [find_positions(e) for e in msg]
        msg_cypher = ''.join([cypher(p) for p in positions])
        return msg_cypher


class TestDigram(unittest.TestCase):
    def setUp(self):
        pass

    def test_message_to_digram_1(self):
        message = "Eu vou ficar bom"
        expected = ['EU', 'VO', 'UF', 'IC', 'AR', 'BO', 'MW']
        self.assertListEqual(message_to_digram(message), expected)

    def test_message_to_digram_2(self):
        message = "O rato roeu "
        expected = ["OR", "AT", "OR", "OE", "UW"]
        self.assertListEqual(message_to_digram(message), expected)

    def test_message_to_digram_eq_letter(self):
        message = "nossa casa"
        expected = ["NO", "SW", "SA", "CA", "SA"]
        self.assertListEqual(message_to_digram(message), expected)


class TestPositions(unittest.TestCase):

    def test_elements_positions_line(self):
        par = 'AC'
        expected = [Position(line=0, column=4), Position(line=1, column=2)]
        self.assertListEqual(find_positions(par), expected)

    def test_elements_positions_Column(self):
        par = 'DT'
        expected = [Position(line=1, column=3), Position(line=3, column=3)]
        self.assertListEqual(find_positions(par), expected)

    def test_elements_positions_diff_line_column(self):
        par = 'AQ'
        expected = [Position(line=0, column=4), Position(line=3, column=1)]
        self.assertListEqual(find_positions(par), expected)


class TestSameLine(unittest.TestCase):
    def test_elements_same_line_AC(self):
        par = 'AE'
        positions = find_positions(par)
        self.assertTrue(same_line(positions))

    def test_elements_same_line_XY(self):
        par = 'XY'
        positions = find_positions(par)
        self.assertTrue(same_line(positions))

    def test_elements_same_line_NT(self):
        par = 'NT'
        positions = find_positions(par)
        self.assertTrue(same_line(positions))

    def test_elements_non_same_line_DT(self):
        par = 'DT'
        positions = find_positions(par)
        self.assertFalse(same_line(positions))


class TestSameColumn(unittest.TestCase):
    def test_elements_same_column_PK(self):
        par = 'PK'
        positions = find_positions(par)
        self.assertTrue(same_column(positions))

    def test_elements_same_column_LY(self):
        par = 'LY'
        positions = find_positions(par)
        self.assertTrue(same_column(positions))

    def test_elements_same_column_IL(self):
        par = 'IL'
        positions = find_positions(par)
        self.assertTrue(same_column(positions))

    def test_elements_non_same_column_IH(self):
        par = 'IH'
        positions = find_positions(par)
        self.assertFalse(same_column(positions))


class TestCypher(unittest.TestCase):
    def setUp(self):
        pass

    def test_cypher_par_sameline_AS(self):
        """ Test if A and S is on same line"""
        par = 'AS'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'EP')

    def test_cypher_par_sameline_SU(self):
        par = 'SI'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'PA')

    def test_cypher_par_sameline_VZ(self):
        par = 'VZ'
        positions = find_positions(par)

        self.assertEqual(cypher(positions), 'WV')

    def test_elements_par_same_line_NT(self):
        par = 'NT'
        positions = find_positions(par)
        self.assertTrue(cypher(positions), 'QU')

    def test_cyper_par_samecolumn_PK(self):
        par = 'PK'
        positions = find_positions(par)

        self.assertEqual(cypher(positions), 'CR')

    def test_cyper_par_samecolumn_IL(self):
        par = 'IL'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'DT')

    def test_cyper_par_samecolumn_IY(self):
        par = 'IY'
        positions = find_positions(par)

        self.assertEqual(cypher(positions), 'DI')

    def test_cyper_par_diff_line_colum_AQ(self):
        """Test cyper digram for different line and column"""
        par = 'AQ'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'SU')

    def test_cyper_par_diff_line_colum_EW(self):
        """Test cyper digram for different line and column"""
        par = 'EW'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'SV')

    def test_cyper_par_diff_line_colum_OW(self):
        """Test cyper digram for different line and column"""
        par = 'OW'
        positions = find_positions(par)
        self.assertEqual(cypher(positions), 'BV')


class TestEncrypt(unittest.TestCase):
    def test_encrypt(self):
        msg = "nossa esperanca e a sua vinda"
        expected = "NIGFAQUFNOQUSAUODTUTIFBV"
        self.assertEqual(encrypt_playfair(msg, act="ENC"), expected)


if __name__ == "__main__":
    unittest.main()