HTB Cyber Apocalypse 2022 - Quantum Engine

crypto quantum
  • Author: Quintec
  • Date:

Let’s take a look at the source code:

  1import socketserver
  2from secrets import flag
  3import signal
  4from qiskit import *
  5import itertools
  6
  7class CircuitException(Exception):
  8    pass
  9
 10
 11class Circuit:
 12
 13    def __init__(self, inputGates) -> None:
 14        self.a = 0
 15        self.b = 0
 16        self.c = 0
 17        self.inputGates = inputGates
 18    
 19    def append_HGate(self, qbit):
 20        qbit = int(qbit)
 21        if qbit in range(4):
 22            self.circuit.h(qbit)
 23        else:
 24            raise CircuitException('Non-valid qbit position given...')
 25    
 26    def append_TGate(self, qbit):
 27        qbit = int(qbit)
 28        if qbit in range(4):
 29            self.circuit.t(qbit)
 30        else:
 31            raise CircuitException('Non-valid qbit position given...')
 32
 33    def append_TdgGate(self, qbit):
 34        qbit = int(qbit)
 35        if qbit in range(4):
 36            self.circuit.tdg(qbit)
 37        else:
 38            raise CircuitException('Non-valid qbit position given...')
 39
 40    def append_CZGate(self, qbits):
 41        qbits = qbits.split(',')
 42        c_qubit = int(qbits[0])
 43        t_qubit = int(qbits[1])
 44
 45        if c_qubit in range(4) and t_qubit in range(4):
 46            self.circuit.cz(c_qubit, t_qubit)
 47        else:
 48            raise CircuitException('Non-valid qbit position given...')
 49
 50
 51    def generate_Circuit(self):
 52
 53        self.circuit = QuantumCircuit(4,3)
 54
 55        if self.a == 1:
 56            self.circuit.x(0)
 57        if self.b == 1:
 58            self.circuit.x(1)
 59        if self.c == 1:
 60            self.circuit.x(2)
 61        
 62        for gate in self.inputGates:
 63            gate = gate.split(':')
 64            if gate[0] == 'H':
 65                self.append_HGate(gate[1])
 66            elif gate[0] == 'T':
 67                self.append_TGate(gate[1])
 68            elif gate[0] == 'TDG':
 69                self.append_TdgGate(gate[1])
 70            elif gate[0] == 'CZ':
 71                self.append_CZGate(gate[1])
 72            else:
 73                raise CircuitException('Non-valid gate given...')
 74
 75        self.circuit.measure([0,2,3],[0,1,2])
 76        if self.circuit.depth() > 43:
 77            raise CircuitException('Circuit is too big...')
 78        
 79    def check_Circuit(self):
 80        inputs = list(itertools.product([0, 1], repeat=3))
 81
 82        for input in inputs:
 83            self.a = input[0]
 84            self.b = input[1]
 85            self.c = input[2]
 86
 87            self.generate_Circuit()
 88
 89            simulator = Aer.get_backend('qasm_simulator')
 90            counts = execute(self.circuit,backend=simulator, shots = 1).result().get_counts()
 91            counts = next(iter(counts))[::-1]
 92            if (int(counts[0]) == self.a) and int(counts[1]) == self.c ^self.a ^ self.b and (int(counts[2]) == self.a&self.b|self.b&self.c|self.c&self.a):
 93                pass
 94            else:
 95                return False
 96        return True
 97
 98        
 99
100def challenge(req):
101    try:
102        req.sendall(b'\lvert--------------------------------\lvert\n'+
103                    b'\lvert Phalcon\'s Accelaration System \lvert\n'+
104                    b'\lvert--------------------------------\lvert\n'+
105                    b'\lvert > Send quantum circuit for the \lvert\n'+
106                    b'\lvert system to analyze...           \lvert\n'+
107                    b'\lvert--------------------------------\lvert\n'+
108                    b'\n> '
109                    )
110        input = req.recv(4096).decode().strip().split(';')
111
112        if len(input) < 0 or len(input) > 100:
113            raise CircuitException('Non-valid circuit length...')
114        
115        quantumCircuit = Circuit(input)
116        
117        if quantumCircuit.check_Circuit():
118            req.sendall(flag.encode()+b"\n")
119            req.close()
120            exit()
121        else:
122            req.sendall(b'The circuit failed to pass the test...\n')
123            req.close()
124            exit()
125
126
127    except CircuitException as ce:
128        try:
129            req.sendall(ce.encode()+b'\n')
130            req.close()
131        except:
132            pass
133        exit()
134
135    except Exception as e:
136        try:
137            req.sendall(b'Unexpected error.\n')
138            req.close()
139        except:
140            pass
141        exit()
142
143class incoming(socketserver.BaseRequestHandler):
144    def handle(self):
145        signal.alarm(300)
146        req = self.request
147        print("starting server")
148        while True:
149            challenge(req)
150
151class ReusableTCPServer(socketserver.ForkingMixIn, socketserver.TCPServer):
152    pass
153
154socketserver.TCPServer.allow_reuse_address = False
155server = ReusableTCPServer(("0.0.0.0", 1337), incoming)
156server.serve_forever()

Alright, so the main part of the challenge is in the Circuit class. It looks like we’re implementing a quantum circuit that satisfies this:

 1def check_Circuit(self):
 2        inputs = list(itertools.product([0, 1], repeat=3))
 3
 4        for input in inputs:
 5            self.a = input[0]
 6            self.b = input[1]
 7            self.c = input[2]
 8
 9            self.generate_Circuit()
10
11            simulator = Aer.get_backend('qasm_simulator')
12            counts = execute(self.circuit,backend=simulator, shots = 1).result().get_counts()
13            counts = next(iter(counts))[::-1]
14            if (int(counts[0]) == self.a) and int(counts[1]) == self.c ^self.a ^ self.b and (int(counts[2]) == self.a&self.b|self.b&self.c|self.c&self.a):
15                pass
16            else:
17                return False
18        return True

So given 3 inputs $a, b, c$ and 4 total qubits to work with, we need to set the outputs to be $a$, $a \oplus b \oplus c$, and $(a\&b)\lvert(b\&c)\lvert(c\&a)$. And it looks like the only gates that we are able to use are $H, T, TDG$, and $CZ$.

Initial Thoughts

Note the line self.circuit.measure([0,2,3],[0,1,2]) in the generate_Circuit function. This means that the outputs need to be on qubits $0, 2, $ and $3$ (and the inputs are on qubits $0, 1, $ and $2$ as shown in the beginning of the function).

The first output is simple by itself - just don’t do anything to the input! This is actually a bit tricky once we implement the others, as you soon will see.

The second output needs to be the XOR of all the inputs. It turns out this isn’t too hard either - a $CNOT$ (or $CX$) gate effectively performs XOR. (Remember that a $CNOT$ gate flips the second qubit (across the $x$ axis, hence the name $CX$) iff the first qubit is $\lvert1\rangle$.) That is, for classical bits, $\text{CNOT}(a, b)$ takes $(a, b)$ to $(a, a \oplus b)$. There’s just one small problem - we aren’t allowed $CNOT$ gates! But we are allowed $CZ$ gates, which flips the second qubit across the $z$ axis iff the first qubit is $\lvert1\rangle$. These two operations only differ by an axis, and it turns out that $HZH = X$ (and $HXH = Z$), so we can implement $CNOT$ gates in this way. Remembering that this output needs to be on qubit 2, we can simply apply $CNOT$s with control qubits 0 and 1 targeting qubit 2, which would take $(a, b, c)$ to $(a, b, a \oplus b \oplus c)$ as desired.

 1qc = QuantumCircuit(4, 3)
 2qc.h(2)
 3qc.cz(0, 2)
 4qc.cz(1, 2)
 5qc.h(2)
 6qc.draw()
 7"""
 8q_0: ──────■─────────
 9           │         
10q_1: ──────┼──■──────
11     ┌───┐ │  │ ┌───┐
12q_2: ┤ H ├─■──■─┤ H ├
13     └───┘      └───┘
14q_3: ────────────────
15                     
16c: 3/════════════════   
17"""

The third output is a bit trickier. $(a\&b)\lvert(b\&c)\lvert(c\&a)$ is a complicated expression. While it is possible to construct circuits that simulate AND and OR, the main problem is that to be able to use the value of each input twice, we would need multiple ancillary qubits so that the value of the inputs aren’t lost. Instead, it helps to think of the expression as a whole rather than the individual parts. What it’s really saying is “set the output bit to $\lvert1\rangle$ if at least 2 of the input bits are $\lvert1\rangle$”. Or, if you think about it, “set the output bit to the majority of the input bits”. What does this sound like? To me, error correction comes to mind.

The basic single-qubit bit flip error correction circuit looks like this:

                ░           ┌───┐
q_0: ──■────■───░───■────■──┤ X ├
     ┌─┴─┐  │   ░ ┌─┴─┐  │  └─┬─┘
q_1: ┤ X ├──┼───░─┤ X ├──┼────■──
     └───┘┌─┴─┐ ░ └───┘┌─┴─┐  │  
q_2: ─────┤ X ├─░──────┤ X ├──■──
          └───┘ ░      └───┘

where $q_0$ is the qubit to be transmitted, and $q_1$ and $q_2$ are ancilla qubits which start in state $\lvert0 \rangle$. The first 2 $CNOT$ gates essentially serve to copy the value of $q_0$ into $q_1$ and $q_2$, and then $q_0$ is potentially flipped. After a few more gates, no matter the starting value of $q_0$ and no matter if it is flipped or not, in the end $q_0$ will have the same value it started with. The last gate is a $CCNOT$ gate, a double-controlled not gate, also known as a Toffoli gate, which flips the target ($q_0$) iff both inputs ($q_1$ and $q_2$) are $\lvert1 \rangle$. It basically acts like an AND gate.

Now it turns out if you disregard the first half of the circuit (copying $q_0$ into $q_1$ and $q_2$), this circuit does exactly what we want it to. Try it yourself for some random inputs if you’re not convinced - it will set $q_0$ to the majority bit. Now how do we implement this? I’ve already shown how to implement $CNOT$ gates above, and for the $CCNOT$ gate, you could either look it up, or use qiskit’s decompose function to decompose it into gates we know how to make.

 1qc = QuantumCircuit(3)
 2qc.ccx(0, 1, 2)
 3qc.decompose().draw()
 4"""
 5                                                       ┌───┐      
 6q_0: ───────────────────■─────────────────────■────■───┤ T ├───■──
 7                        │             ┌───┐   │  ┌─┴─┐┌┴───┴┐┌─┴─┐
 8q_1: ───────■───────────┼─────────■───┤ T ├───┼──┤ X ├┤ TDG ├┤ X ├
 9     ┌───┐┌─┴─┐┌─────┐┌─┴─┐┌───┐┌─┴─┐┌┴───┴┐┌─┴─┐├───┤└┬───┬┘└───┘
10q_2: ┤ H ├┤ X ├┤ TDG ├┤ X ├┤ T ├┤ X ├┤ TDG ├┤ X ├┤ T ├─┤ H ├──────
11     └───┘└───┘└─────┘└───┘└───┘└───┘└─────┘└───┘└───┘ └───┘      
12"""

Putting it All Together

There are a few final issues to take care of. We have successfully calculated each output bit individually, but how to put it all together? First, note that our adapted error correction circuit has the potential to flip $q_1$ and $q_2$, so we need to track the original value of $q_0$ to see whether or not we need to re-flip $q_1$ and $q_2$ after performing the “error correction”. This is easy enough - since our extra 4th qubit $q_3$ is guaranteed to start at zero, performing a $CNOT$ to XOR the value of one of the qubits into the 4th qubit would copy it into the 4th qubit. We can actually then just run our adapted error correction circuit targeting this 4th qubit, since we want this output on $q_3$ in the end anyway. After that, since we have the original value of the qubit, we can reflip the other qubits if needed, and then use all 3 original inputs to calculate the other outputs, which is not hard. The final circuit looks like this:

  1from qiskit import *
  2from qiskit.extensions import *
  3
  4qc = QuantumCircuit(4, 3)
  5
  6#copy qubit 3 to qubit 2
  7qc.h(3)
  8qc.cz(2, 3)
  9qc.h(3)
 10
 11#"error correct" (output 3): convert qubit 3 to majority of (0, 1, 3), using two CX and one CCX 
 12qc.h(0)
 13qc.cz(3, 0)
 14qc.h(0)
 15qc.h(1)
 16qc.cz(3, 1)
 17qc.h(1)
 18
 19#CCX decomposed
 20qc.cz(1, 3)
 21qc.h(3)
 22qc.tdg(3)
 23qc.h(3)
 24qc.cz(0, 3)
 25qc.h(3)
 26qc.t(3)
 27qc.h(3)
 28qc.cz(1, 3)
 29qc.h(3)
 30qc.tdg(3)
 31qc.h(3)
 32qc.cz(0, 3)
 33qc.h(3)
 34qc.t(3)
 35qc.h(3)
 36qc.t(1)
 37qc.h(1)
 38qc.cz(0, 1)
 39qc.h(1)
 40qc.tdg(1)
 41qc.t(0)
 42qc.h(1)
 43qc.cz(0, 1)
 44qc.h(1)
 45
 46#reflip 0 and 1 if necessary
 47qc.h(0)
 48qc.cz(2, 0)
 49qc.h(0)
 50qc.h(1)
 51qc.cz(2, 1)
 52qc.h(1)
 53
 54#xor (output 2)
 55qc.h(2)
 56qc.cz(0, 2)
 57qc.cz(1, 2)
 58qc.h(2)
 59
 60qc.measure([0,2,3],[0,1,2])
 61
 62qc.draw()
 63"""
 64     ┌───┐           ┌───┐                    »
 65q_0: ┤ H ├─────────■─┤ H ├────────────────────»
 66     ├───┤         │ └───┘┌───┐               »
 67q_1: ┤ H ├─────────┼───■──┤ H ├─■─────────────»
 68     └───┘         │   │  └───┘ │             »
 69q_2: ──────■───────┼───┼────────┼─────────────»
 70     ┌───┐ │ ┌───┐ │   │        │ ┌───┐┌─────┐»
 71q_3: ┤ H ├─■─┤ H ├─■───■────────■─┤ H ├┤ TDG ├»
 72     └───┘   └───┘                └───┘└─────┘»
 73c: 3/═════════════════════════════════════════»
 74                                              »
 75«                                                »
 76«q_0: ──────■────────────────────────────────────»
 77«           │                   ┌───┐ ┌───┐      »
 78«q_1: ──────┼─────────────────■─┤ T ├─┤ H ├──────»
 79«           │                 │ └───┘ └───┘      »
 80«q_2: ──────┼─────────────────┼──────────────────»
 81«     ┌───┐ │ ┌───┐┌───┐┌───┐ │ ┌───┐┌─────┐┌───┐»
 82«q_3: ┤ H ├─■─┤ H ├┤ T ├┤ H ├─■─┤ H ├┤ TDG ├┤ H ├»
 83«     └───┘   └───┘└───┘└───┘   └───┘└─────┘└───┘»
 84«c: 3/═══════════════════════════════════════════»
 85«                                                »
 86«             ┌───┐               ┌───┐   ┌───┐»
 87«q_0: ─■───■──┤ T ├─────────────■─┤ H ├─■─┤ H ├»
 88«      │   │  ├───┤┌─────┐┌───┐ │ ├───┤ │ ├───┤»
 89«q_1: ─┼───■──┤ H ├┤ TDG ├┤ H ├─■─┤ H ├─┼─┤ H ├»
 90«      │      └───┘└─────┘└───┘   └───┘ │ └───┘»
 91«q_2: ─┼────────────────────────────────■──────»
 92«      │ ┌───┐┌───┐ ┌───┐  ┌─┐                 »
 93«q_3: ─■─┤ H ├┤ T ├─┤ H ├──┤M├─────────────────»
 94«        └───┘└───┘ └───┘  └╥┘                 »
 95«c: 3/══════════════════════╩══════════════════»
 96«                           2                  »
 97«                   ┌─┐        
 98«q_0: ─────────■────┤M├────────
 99«        ┌───┐ │    └╥┘        
100«q_1: ─■─┤ H ├─┼──■──╫─────────
101«      │ ├───┤ │  │  ║ ┌───┐┌─┐
102«q_2: ─■─┤ H ├─■──■──╫─┤ H ├┤M├
103«        └───┘       ║ └───┘└╥┘
104«q_3: ───────────────╫───────╫─
105«                    ║       ║ 
106«c: 3/═══════════════╩═══════╩═
107«                    0       1 
108"""

Converting this to the input format given in the source code and submitting it to the server gives us the flag.