CTF-2022/crypto/pai3/paiaiai.py at main · hackpack-ncsu/CTF-2022 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#!/usr/bin/env python3
#
# Polymero
#

# Imports
from Crypto.Util.number import getPrime, inverse
from secrets import randbelow

# Local imports
with open("flag.txt",'rb') as f:
    FLAG = f.read().decode()
    f.close()

# Just for you sanity
assert len(FLAG) > 64


HDR = r"""|
|   __________     ___   _____  .___  ___     /\  ________
|   \______   \   /  /  /  _  \ |   | \  \   /  \ \_____  \
|    |     ___/  /  /  /  /_\  \|   |  \  \  \/\/   _(__  <
|    |    |     (  (  /    |    \   |   )  )       /       \
|    |____|      \  \ \____|__  /___|  /  /       /______  /
|                 \__\        \/      /__/               \/
|"""

MENU = r"""|
|  MENU:
|   [E]ncrypt
|   [D]ecrypt
|   [Q]uit
|"""


class Paiaiai:
    """ My first Paillier implementation! So proud of it. ^ w ^ """

    def __init__(self):
        # Key generation
        p, q = [getPrime(512) for _ in range(2)]
        n = p * q
        # Public key
        self.pub = {
            'n'  : n,
            'gp' : pow(randbelow(n**2), p, n**2),
            'gq' : pow(randbelow(n**2), q, n**2)
        }
        # Private key
        self.priv = {
            'la' : (p - 1)*(q - 1),
            'mu' : inverse((pow(self.pub['gp'] * self.pub['gq'], (p-1)*(q-1), n**2) - 1) // n, n)
        }

    def encrypt(self, m: str):
        m_int = int.from_bytes(m.encode(), 'big')
        g = pow([self.pub['gp'],self.pub['gq']][randbelow(2)], m_int, self.pub['n']**2)
        r = pow(randbelow(self.pub['n']), self.pub['n'], self.pub['n']**2)
        return (g * r) % self.pub['n']**2

    def decrypt(self, c: int):
        cl = (pow(c, self.priv['la'], self.pub['n']**2) - 1) // self.pub['n']
        return (cl * self.priv['mu']) % self.pub['n']


print(HDR)
pai = Paiaiai()

while True:

    try:

        print(MENU)
        choice = input("|  >>> ").lower().strip()

        if choice == 'e':
            print("|\n|  ENCRYPT:")
            print("|   [F]lag")
            print("|   [M]essage")
            subchoice = input("|\n|  >>> ").lower().strip()

            if subchoice == 'f':
                enc_flag = pai.encrypt(FLAG)
                print("|\n|  FLAG:", enc_flag)

            elif subchoice == 'm':
                msg = input("|\n|  MSG: str\n|   > ")
                cip = pai.encrypt(msg)
                print("|\n|  CIP:", cip)

        elif choice == 'd':
            cip = input("|\n|  CIP: int\n|   > ")
            msg = pai.decrypt(int(cip))
            print("|\n|  MSG:", msg)

        elif choice == 'q':
            print("|\n|  Bai ~ \n|")
            break

        else:
            print("|\n|  Trai again ~ \n|")

    except (KeyboardInterrupt, EOFError):
        print("\n|\n|  Bai ~ \n|")
        break

    except:
        print("|\n|  Aiaiai ~ \n|")