File size: 10,927 Bytes
6142a25
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
import bchlib
import numpy as np 
from typing import List, Tuple
import random 
from copy import deepcopy

class RSC(object):
    def __init__(self, data_bytes=16, ecc_bytes=4, verbose=False, **kwargs):
        from reedsolo import RSCodec
        self.rs = RSCodec(ecc_bytes)
        if verbose:
            print(f'Reed-Solomon ECC len: {ecc_bytes*8} bits')
        self.data_len = data_bytes
        self.dlen = data_bytes * 8  # data length in bits
        self.ecc_len = ecc_bytes * 8  # ecc length in bits
    
    def get_total_len(self):
        return self.dlen + self.ecc_len

    def encode_text(self, text: List[str]):
        return np.array([self._encode_text(t) for t in text])
    
    def _encode_text(self, text: str):
        text = text + ' ' * (self.dlen // 8 - len(text))
        out = self.rs.encode(text.encode('utf-8'))  # bytearray
        out = ''.join(format(x, '08b') for x in out)  # bit string
        out = np.array([int(x) for x in out], dtype=np.float32)
        return out
    
    def decode_text(self, data: np.array):
        assert len(data.shape)==2
        return [self._decode_text(d) for d in data]

    def _decode_text(self, data: np.array):
        assert len(data.shape)==1
        data = ''.join([str(int(bit)) for bit in data])
        data = bytes(int(data[i: i + 8], 2) for i in range(0, len(data), 8))
        data = bytearray(data)
        try:
            data = self.rs.decode(data)[0]
            data = data.decode('utf-8').strip()
        except:
            print('Error: Decode failed')
            data = get_random_unicode(self.get_total_len()//8)
        
        return data

def get_random_unicode(length):
    # Update this to include code point ranges to be sampled
    include_ranges = [
        ( 0x0021, 0x0021 ),
        ( 0x0023, 0x0026 ),
        ( 0x0028, 0x007E ),
        ( 0x00A1, 0x00AC ),
        ( 0x00AE, 0x00FF ),
        ( 0x0100, 0x017F ),
        ( 0x0180, 0x024F ),
        ( 0x2C60, 0x2C7F ),
        ( 0x16A0, 0x16F0 ),
        ( 0x0370, 0x0377 ),
        ( 0x037A, 0x037E ),
        ( 0x0384, 0x038A ),
        ( 0x038C, 0x038C ),
    ]
    alphabet = [
        chr(code_point) for current_range in include_ranges
            for code_point in range(current_range[0], current_range[1] + 1)
    ]
    return ''.join(random.choice(alphabet) for i in range(length))


class BCH(object):
    def __init__(self, BCH_POLYNOMIAL = 137, BCH_BITS = 5, payload_len=100, verbose=True,**kwargs):
        self.bch = bchlib.BCH(BCH_POLYNOMIAL, BCH_BITS)
        self.payload_len = payload_len  # in bits
        self.data_len = (self.payload_len - self.bch.ecc_bytes*8)//7  # in ascii characters
        assert self.data_len*7+self.bch.ecc_bytes*8 <= self.bch.n, f'Error! BCH with poly {BCH_POLYNOMIAL} and bits {BCH_BITS} can only encode max {self.bch.n//8} bytes of total payload'
        if verbose:
            print(f'BCH: POLYNOMIAL={BCH_POLYNOMIAL}, protected bits={BCH_BITS}, payload_len={payload_len} bits, data_len={self.data_len*7} bits ({self.data_len} ascii chars), ecc len={self.bch.ecc_bytes*8} bits')
    
    def get_total_len(self):
        return self.payload_len
    
    def encode_text(self, text: List[str]):
        return np.array([self._encode_text(t) for t in text])
    
    def _encode_text(self, text: str):
        text = text + ' ' * (self.data_len - len(text))
        # data = text.encode('utf-8')  # bytearray
        data = encode_text_ascii(text)  # bytearray
        ecc = self.bch.encode(data)  # bytearray
        packet = data + ecc  # payload in bytearray
        packet = ''.join(format(x, '08b') for x in packet)
        packet = [int(x) for x in packet]
        packet.extend([0]*(self.payload_len - len(packet)))
        packet = np.array(packet, dtype=np.float32)
        return packet
    
    def decode_text(self, data: np.array):
        assert len(data.shape)==2
        return [self._decode_text(d) for d in data]
    
    def _decode_text(self, packet: np.array):
        assert len(packet.shape)==1
        packet = ''.join([str(int(bit)) for bit in packet])  # bit string
        packet = packet[:(len(packet)//8*8)]  # trim to multiple of 8 bits
        packet = bytes(int(packet[i: i + 8], 2) for i in range(0, len(packet), 8))
        packet = bytearray(packet)
        # assert len(packet) == self.data_len + self.bch.ecc_bytes
        data, ecc = packet[:-self.bch.ecc_bytes], packet[-self.bch.ecc_bytes:]
        data0 = decode_text_ascii(deepcopy(data)).strip()
        bitflips = self.bch.decode_inplace(data, ecc)
        if bitflips == -1:  # error, return random text
            data = data0
        else:
            # data = data.decode('utf-8').strip()
            data = decode_text_ascii(data).strip()
        return data


def encode_text_ascii(text: str):
    # encode text to 7-bit ascii
    # input: text, str
    # output: encoded text, bytearray
    text_int7 = [ord(t) & 127 for t in text]
    text_bitstr = ''.join(format(t,'07b') for t in text_int7)
    if len(text_bitstr) % 8 != 0:
        text_bitstr =  '0'*(8-len(text_bitstr)%8) + text_bitstr  # pad to multiple of 8
    text_int8 = [int(text_bitstr[i:i+8], 2) for i in range(0, len(text_bitstr), 8)]
    return bytearray(text_int8)


def decode_text_ascii(text: bytearray):
    # decode text from 7-bit ascii
    # input: text, bytearray
    # output: decoded text, str
    text_bitstr = ''.join(format(t,'08b') for t in text)  # bit string
    pad = len(text_bitstr) % 7
    if pad != 0:  # has padding, remove
        text_bitstr = text_bitstr[pad:]
    text_int7 = [int(text_bitstr[i:i+7], 2) for i in range(0, len(text_bitstr), 7)]
    text_bytes = bytes(text_int7)
    return text_bytes.decode('utf-8')


class ECC(object):
    def __init__(self, BCH_POLYNOMIAL = 137, BCH_BITS = 5, **kwargs):
        self.bch = bchlib.BCH(BCH_POLYNOMIAL, BCH_BITS)
    
    def get_total_len(self):
        return 100

    def _encode(self, x):
        # x: 56 bits, {0, 1}, np.array
        # return: 100 bits, {0, 1}, np.array
        dlen = len(x)
        data_str = ''.join(str(x) for x in x.astype(int))
        packet = bytes(int(data_str[i: i + 8], 2) for i in range(0, dlen, 8))
        packet = bytearray(packet)
        ecc = self.bch.encode(packet)
        packet = packet + ecc  # 96 bits
        packet = ''.join(format(x, '08b') for x in packet)
        packet = [int(x) for x in packet]
        packet.extend([0, 0, 0, 0])
        packet = np.array(packet, dtype=np.float32)  # 100
        return packet

    def _decode(self, x):
        # x: 100 bits, {0, 1}, np.array
        # return: 56 bits, {0, 1}, np.array
        packet_binary = "".join([str(int(bit)) for bit in x])
        packet = bytes(int(packet_binary[i: i + 8], 2) for i in range(0, len(packet_binary), 8))
        packet = bytearray(packet)

        data, ecc = packet[:-self.bch.ecc_bytes], packet[-self.bch.ecc_bytes:]
        bitflips = self.bch.decode_inplace(data, ecc)
        if bitflips == -1:  # error, return random data
            data = np.random.binomial(1, .5, 56) 
        else:
            data = ''.join(format(x, '08b') for x in data)
            data = np.array([int(x) for x in data], dtype=np.float32)
        return data  # 56 bits
    
    def _generate(self):
        dlen = 56
        data= np.random.binomial(1, .5, dlen)
        packet = self._encode(data)
        return packet, data

    def generate(self, nsamples=1):
        # generate random 56 bit secret
        data = [self._generate() for _ in range(nsamples)]
        data = (np.array([d[0] for d in data]), np.array([d[1] for d in data]))
        return data  # data with ecc, data org
    
    def _to_text(self, data):
        # data:  {0, 1}, np.array
        # return: str
        data = ''.join([str(int(bit)) for bit in data])
        all_bytes = [ data[i: i+8] for i in range(0, len(data), 8) ]
        text = ''.join([chr(int(byte, 2)) for byte in all_bytes])
        return text.strip()
    
    def _to_binary(self, s):
        if isinstance(s, str):
            out = ''.join([ format(ord(i), "08b") for i in s ])
        elif isinstance(s, bytes):
            out = ''.join([ format(i, "08b") for i in s ])
        elif isinstance(s, np.ndarray) and s.dtype is np.dtype(bool):
            out = ''.join([chr(int(i)) for i in s])
        elif isinstance(s, int) or isinstance(s, np.uint8):
            out = format(s, "08b")
        elif isinstance(s, np.ndarray):
            out = [ format(i, "08b") for i in s ]
        else:
            raise TypeError("Type not supported.")

        return np.array([float(i) for i in out], dtype=np.float32)

    def _encode_text(self, s):
        s = s + ' '*(7-len(s))  # 7 chars
        s = self._to_binary(s)  # 56 bits
        packet = self._encode(s)  # 100 bits
        return packet, s

    def encode_text(self, secret_list, return_pre_ecc=False):
        """encode secret with BCH ECC.
        Input: secret (list of strings)
        Output: secret (np array) with shape (B, 100) type float23, val {0,1}"""
        assert np.all(np.array([len(s) for s in secret_list]) <= 7), 'Error! all strings must be less than 7 characters'
        secret_list = [self._encode_text(s) for s in secret_list]
        ecc = np.array([s[0] for s in secret_list], dtype=np.float32)
        if return_pre_ecc:
            return ecc, np.array([s[1] for s in secret_list], dtype=np.float32)
        return ecc
    
    def decode_text(self, data):
        """Decode secret with BCH ECC and convert to string.
        Input: secret (torch.tensor) with shape (B, 100) type bool
        Output: secret (B, 56)"""
        data = self.decode(data)
        data = [self._to_text(d) for d in data]
        return data

    def decode(self, data):
        """Decode secret with BCH ECC and convert to string.
        Input: secret (torch.tensor) with shape (B, 100) type bool
        Output: secret (B, 56)"""
        data = data[:, :96]
        data = [self._decode(d) for d in data]
        return np.array(data)

def test_ecc():
    ecc = ECC()
    batch_size = 10 
    secret_ecc, secret_org = ecc.generate(batch_size)  # 10x100 ecc secret, 10x56 org secret
    # modify secret_ecc
    secret_pred = secret_ecc.copy()
    secret_pred[:,3:6] = 1 - secret_pred[:,3:6]
    # pass secret_ecc to model and get predicted as secret_pred
    secret_pred_org = ecc.decode(secret_pred)  # 10x56
    assert np.all(secret_pred_org == secret_org)  # 10


def test_bch():
    # test 100 bit
    def check(text, poly, k, l):
        bch = BCH(poly, k, l)
        # text = 'secrets'
        encode = bch.encode_text([text])
        for ind in np.random.choice(l, k):
            encode[0, ind] = 1 - encode[0, ind]
        text_recon = bch.decode_text(encode)[0]
        assert text==text_recon
    
    check('secrets', 137, 5, 100)
    check('some secret', 285, 10, 160)

if __name__ == '__main__':
    test_ecc()
    test_bch()