<commit_before>#!/usr/bin/python

import argparse
import base64
import json
from suffix_tree import SuffixTree
from Crypto.Cipher import AES
import hashlib

parser = argparse.ArgumentParser()
parser.add_argument('--filename', '-f', default='../files/small_shakespeare.txt',
                    help='Path of file to encrypt')
args = parser.parse_args()


from Crypto.Cipher import AES
from Crypto import Random

BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS) 
unpad = lambda s : s[0:-ord(s[-1])]

K_1 = 'This is a key789'
K_D = 'This is a key123'
K_C = 'This is a key234'
IV_D = 'This is an IV456'
IV_C = 'This is an IV567'

K_1 += K_1
K_D += K_D
K_C += K_C

filename = args.filename

# key = 'This is a key12 This is a key12 '
# text = 'string of size16'
# iv = 'This is an IV456'
# aes = AES.new(key, AES.MODE_CBC, iv)
# ciphertext = aes.encrypt(pad(text))
# ciphertext = base64.b64encode(ciphertext)
#
# aes = AES.new(key, AES.MODE_CBC, iv)
# plaintext = aes.decrypt(base64.b64decode(ciphertext))
# print unpad(plaintext)

st = None
length = None
with open(filename) as f:
  st = SuffixTree(f.read())
  length = len(st.string) - 1

D = {}
for leaf in st.leaves:
  leafPL = leaf.pathLabel[:len(leaf.pathLabel)-1] + '$'
  parentPL = leaf.parent.pathLabel
  initPath = leafPL[:len(parentPL)+1]
  key = hashlib.sha256(initPath.encode()).hexdigest()
  index = '{:>16}'.format(str(length - len(leaf.pathLabel)))
  aes_D = AES.new(K_D, AES.MODE_CBC, IV_D)
  raw = aes_D.encrypt(pad(index))
  value = base64.b64encode(raw)
  D[key] = value

C = [None] * (len(st.string) - 1)
for i, c in enumerate(st.string[:-1]):
  aes_C = AES.new(K_C, AES.MODE_CBC, IV_C)
  raw = aes_C.encrypt(pad('{:>16}'.format(c)))
  value = base64.b64encode(raw)
  C[i] = value

with open('ciphertext.txt', 'w') as fout:<commit_msg>fix tree bug; encrypt inner nodes<commit_after>#!/usr/bin/python

import argparse
import base64
import json
from suffix_tree import SuffixTree
from Crypto.Cipher import AES
import hashlib

parser = argparse.ArgumentParser()
parser.add_argument('--filename', '-f', default='../files/small_shakespeare.txt',
                    help='Path of file to encrypt')
args = parser.parse_args()


from Crypto.Cipher import AES
from Crypto import Random

BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS) 
unpad = lambda s : s[0:-ord(s[-1])]

K_1 = 'This is a key789'
K_D = 'This is a key123'
K_C = 'This is a key234'
IV_D = 'This is an IV456'
IV_C = 'This is an IV567'

K_1 += K_1
K_D += K_D
K_C += K_C

filename = args.filename

# key = 'This is a key12 This is a key12 '
# text = 'string of size16'
# iv = 'This is an IV456'
# aes = AES.new(key, AES.MODE_CBC, iv)
# ciphertext = aes.encrypt(pad(text))
# ciphertext = base64.b64encode(ciphertext)
#
# aes = AES.new(key, AES.MODE_CBC, iv)
# plaintext = aes.decrypt(base64.b64decode(ciphertext))
# print unpad(plaintext)

st = None
length = None
with open(filename) as f:
  st = SuffixTree(f.read())
  length = len(st.string) - 1

D = {}
for leaf in st.leaves:
  leafPL = leaf.pathLabel[:len(leaf.pathLabel)-1] + '$'
  parentPL = leaf.parent.pathLabel
  initPath = leafPL[:len(parentPL)+1]
  key = hashlib.sha256(initPath.encode()).hexdigest()
  index = '{:>16}'.format(str(length - len(leaf.pathLabel)))
  aes_D = AES.new(K_D, AES.MODE_CBC, IV_D)
  raw = aes_D.encrypt(pad(index))
  value = base64.b64encode(raw)
  D[key] = value

for innerNode in st.innerNodes:
  if innerNode.stringDepth == 0:
    continue
  nodePL = innerNode.pathLabel
  parentPL = innerNode.parent.pathLabel
  initPath = nodePL[:len(parentPL)+1]
  key = hashlib.sha256(initPath.encode()).hexdigest()
  index = '{:>16}'.format(str(length - len(innerNode.pathLabel)))
  aes_D = AES.new(K_D, AES.MODE_CBC, IV_D)
  raw = aes_D.encrypt(pad(index))
  value = base64.b64encode(raw)
  D[key] = value

C = [None] * (len(st.string) - 1)
for i, c in enumerate(st.string[:-1]):
  aes_C = AES.new(K_C, AES.MODE_CBC, IV_C)
  raw = aes_C.encrypt(pad('{:>16}'.format(c)))
  value = base64.b64encode(raw)
  C[i] = value

with open('ciphertext.txt', 'w') as fout:
{'fraction_non_alphanumeric': 0.15821550463188688, 'fraction_numerical': 0.032423208191126277, 'mean_word_length': 3.967312348668281, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1021078', 'n_tokens_mistral': 1762, 'n_tokens_neox': 1631, 'n_words': 399}