49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
from os.path import isdir, join
|
|
from os import listdir
|
|
from db import CwmDb
|
|
from zipfile import ZipFile
|
|
from base64 import b64decode
|
|
|
|
|
|
def import_keys(key: str, db: CwmDb, force=False):
|
|
is_zip = False
|
|
file_list = []
|
|
contain_dir_name = False
|
|
if isdir(key):
|
|
file_list = listdir(key)
|
|
else:
|
|
z = ZipFile(key)
|
|
is_zip = True
|
|
file_list = z.namelist()
|
|
if 'Y2hlcy8/' in file_list:
|
|
contain_dir_name = True
|
|
file_list = [i[8:] for i in file_list if i.startswith(
|
|
'Y2hlcy8/') and i != 'Y2hlcy8/']
|
|
try:
|
|
count = 0
|
|
keys = db.get_all_keys_as_origin()
|
|
for i in file_list:
|
|
oid = b64decode(i).decode()
|
|
if oid in keys:
|
|
if not force:
|
|
continue
|
|
cid = int(oid[0:9])
|
|
uid = int(oid[9:])
|
|
if is_zip:
|
|
if contain_dir_name:
|
|
path = 'Y2hlcy8/' + i
|
|
else:
|
|
path = i
|
|
content = z.read(path).decode()
|
|
else:
|
|
content = open(join(key, i), 'r', encoding='UTF-8').read()
|
|
if oid in keys and content == keys[oid]:
|
|
continue
|
|
db.add_key(cid, uid, content)
|
|
count += 1
|
|
print(f'Imported {count} keys.')
|
|
finally:
|
|
db._db.commit()
|
|
if is_zip:
|
|
z.close()
|