def parse_args():
opts, args = getopt.getopt(sys.argv[1:], '-s-p:-i-g-e-u-f:',
['status', 'project', 'init', 'gen', 'enc', 'upload', 'find='])
client = None
# 先初始化SSEClient对象
for opt_name, opt_val in opts:
if opt_name in ('-p', '--project'):
client = SSEClient(opt_val, 256, 16)
break
if client is None:
printer.print_error('请指定项目名! (使用参数-p <项目名> 或 --project <项目名>)')
sys.exit(1)
for opt_name, opt_val in opts:
if opt_name in ('-s', '--status'):
client.status()
if opt_name in ('-i', '--init'):
client.init()
if opt_name in ('-g', '--gen'):
client.generate_keys()
if opt_name in ('-e', '--enc'):
client.encrypt()
if opt_name in ('-u', '--upload'):
client.upload()
if opt_name in ('-f', '--find'):
keyword = opt_val
if keyword != '':
client.find(keyword)
else:
search_multiple_times(client)
def find(self, keyword):
"""
执行搜索操作
:param keyword: 待搜索的关键词
:return:
"""
def save_result():
try:
if not os.path.isdir(self.proj_dir_path + 'search_results'):
os.mkdir(self.proj_dir_path + 'search_results')
DIR = self.proj_dir_path + 'search_results/' + keyword
if not os.path.isdir(DIR):
os.mkdir(DIR)
DIR += '/'
for index in res:
tmp = download_manager.download_from_server_given_client_and_decrypt(self, index)
printer.print_info(tmp[0]) # 打印出相应的标题
with open(DIR + tmp[0], 'w', encoding='UTF-8') as f:
f.write(tmp[1])
printer.print_info('搜索结果已经保存在目录 ' + DIR + ' 中。')
except OSError:
printer.print_error('由于系统限制，指定的关键词无法搜索!')
def search_action():
return search_manager.search_once_from_server(self, keyword)
if get_status_by_bits(self.status_bits) != 1:
printer.print_error('操作失败，理由: ')
self.status()
return
res = search_action()
printer.print_info('搜索结果如下:')
save_result()
def generate_keys(self):
"""
生成密钥(k1,k2,k3,k4)
:return:
"""
def generate_keys_action():
k1, k2, k3, k4 = self.gen()
print('========THE KEY========')
print('{}\n{}\n{}\n{}'.format(base64.b64encode(k1).decode(encoding='UTF-8'),
base64.b64encode(k2).decode(encoding='UTF-8'),
base64.b64encode(k3).decode(encoding='UTF-8'),
base64.b64encode(k4).decode(encoding='UTF-8')))
print('========THE KEY========')
# 保存密钥
self.save_keys()
printer.print_success('密钥文件已保存至本地.')
if 2 <= get_status_by_bits(self.status_bits) < 4 or get_status_by_bits(self.status_bits) == 7: # 前面的步骤没有完成时
printer.print_error('操作失败，理由: ')
self.status()
return
if get_status_by_bits(self.status_bits) == 6 or get_status_by_bits(self.status_bits) == 1:
# 如果之前已经有了密文集或者已经上传到了服务器
# 需要告知用户谨慎生成密钥文件
printer.print_warning('已发现使用旧密钥加密后的密文集和索引，重新生成密钥需要重新自行执行enc和upload方法进行同步更新.\n'
'是否需要继续? (Y/N)')
ok = input()
if ok == 'Y' or ok == 'y':
generate_keys_action()
else:
printer.print_info('程序没有进行任何操作，退出...')
generate_keys_action()
def __init__(self, proj_name, k=0, l=0):
self.proj_name = proj_name
self.proj_dir_path = proj_name + '/'
if os.path.isdir(self.proj_dir_path): # 如果项目已经存在，不需要指定k和l，直接读取参数文件
self.k, self.l = self.load_config()[:2]
else:
# 检查参数，进行必要的错误或警告报告
if k != 128 and k != 192 and k != 256:
printer.print_error('The key length of AES must be 128, 192 or 256 bits.')
sys.exit(1)
if l % 8 != 0:
printer.print_warning('The length of the parameter l is not an integer multiple of 8.')
self.k = k
self.k = byte_alignment(self.k)
self.l = l # 在论文中，参数l需要指定
self.l = byte_alignment(self.l)
self.s = scanner.get_s(self.proj_dir_path) # todo
self.s = byte_alignment(self.s)
# 状态信息变量
self.exists_plain_texts = False # 是否存在明文集
self.exists_cipher_texts = False # 是否存在密文集
self.exists_key_file = False # 是否存在密钥文件
self.exists_proj_dir = False # 是否存在项目文件夹
# 第一位 --> 是否存在项目文件夹
# 第二位 --> 是否存在密钥文件
# 第三位 --> 是否存在密文集
# 第四位 --> 是否存在明文集
self.status_bits = 0
self.set_status_bits() # 判断当前proj的状态
self.T = [None] * (2 ** self.l)
self.k1, self.k2, self.k3, self.k4 = None, None, None, None
self.load_keys()
# self.D = None
self.distinct_word_set = None
self.D_ = None
# EncK(D) step3. initialize a global counter ctr = 1
self.ctr = 1
self.A = [None] * (2 ** self.s)
self.entry_size_of_A = -1
self.addrA = {} # 组织成dict结构，用于获取每个链表第一个节点对应于A的位置
self.k0_for_each_keyword = {}
self.file_cnt = scanner.get_file_count(self.proj_dir_path)
if self.file_cnt == 0:
self.file_cnt_byte = 0
else:
self.file_cnt_byte = int(math.ceil(math.log2(self.file_cnt) / 8))
def _load_data(self, train_data_path):
if os.path.exists(train_data_path):
self.train_data_df_top_comments = pd.read_csv(train_data_path)[[
'timestamp', 'author_id', 'comment_id', 'article_id',
'parent_comment_id'
]]
printer.print_success('Input Data loaded')
else:
printer.print_error('Train data path does not exist!')
def _load_data(self):
self.article_dates_df = pd.read_csv(self.article_dates_path)
self.article_dates_df['timestamp'] = pd.to_datetime(
self.article_dates_df['date'])
if os.path.exists(self.data_path):
self.data_df = pd.read_csv(self.data_path)[[
'timestamp', 'author_id', 'comment_id', 'article_id',
'parent_comment_id'
]]
printer.print_success('Input Data loaded ')
else:
printer.print_error('Train data path does not exist!')
def encrypt_action():
printer.print_info('检查明文目录下文件名格式是否符合要求...')
if not scanner.check_filename_format(self.proj_dir_path):
printer.print_info('不符合文件命名格式，请问是否需要执行自动格式化文件名操作? (Y/N)')
ok = input()
if ok == 'y' or ok == 'Y':
scanner.reformat_filename(self.proj_dir_path)
printer.print_success('格式化文件名成功!')
else:
printer.print_error('软件终止...请自行更改文件名以满足要求!')
else:
printer.print_success('检查完毕，文件名符合要求!')
printer.print_info('开始加密索引和文档...')
self.enc()
self.save_encrypted_index() # 记得保存索引
printer.print_success('加密索引和文档成功')
def save_result():
try:
if not os.path.isdir(self.proj_dir_path + 'search_results'):
os.mkdir(self.proj_dir_path + 'search_results')
DIR = self.proj_dir_path + 'search_results/' + keyword
if not os.path.isdir(DIR):
os.mkdir(DIR)
DIR += '/'
for index in res:
tmp = download_manager.download_from_server_given_client_and_decrypt(self, index)
printer.print_info(tmp[0]) # 打印出相应的标题
with open(DIR + tmp[0], 'w', encoding='UTF-8') as f:
f.write(tmp[1])
printer.print_info('搜索结果已经保存在目录 ' + DIR + ' 中。')
except OSError:
printer.print_error('由于系统限制，指定的关键词无法搜索!')
def upload(self):
"""
上传密文、加密后的索引、配置文件到服务器上
:return:
"""
def upload_action():
return upload_manager.upload_to_server(self.proj_name, 'Y')
def delete_local_cipher():
"""
如果上传完毕，删除本地上的密文集和加密索引
:return:
"""
shutil.rmtree(self.proj_dir_path + 'cipher_text')
os.remove(self.proj_dir_path + 'index.enc')
if 2 <= get_status_by_bits(self.status_bits) < 6:
printer.print_error('操作失败，理由: ')
self.status()
return
res = upload_action()
if res != 'success':
printer.print_error('上传失败！服务器返回信息如下：')
printer.print_error(res)
else:
printer.print_success('上传成功!')
delete_local_cipher()
def encrypt(self):
"""
生成索引、加密索引和加密文档
:return:
"""
def encrypt_action():
printer.print_info('检查明文目录下文件名格式是否符合要求...')
if not scanner.check_filename_format(self.proj_dir_path):
printer.print_info('不符合文件命名格式，请问是否需要执行自动格式化文件名操作? (Y/N)')
ok = input()
if ok == 'y' or ok == 'Y':
scanner.reformat_filename(self.proj_dir_path)
printer.print_success('格式化文件名成功!')
else:
printer.print_error('软件终止...请自行更改文件名以满足要求!')
else:
printer.print_success('检查完毕，文件名符合要求!')
printer.print_info('开始加密索引和文档...')
self.enc()
self.save_encrypted_index() # 记得保存索引
printer.print_success('加密索引和文档成功')
def delete_local_plain_texts():
"""
加密操作成功后，删除本地上的明文文件
:return:
"""
shutil.rmtree(self.proj_dir_path + 'plain_text')
if 2 <= get_status_by_bits(self.status_bits) < 5 or get_status_by_bits(self.status_bits) == 7:
printer.print_error('操作失败，理由: ')
self.status()
return
encrypt_action()
delete_local_plain_texts()
