156 lines
4.6 KiB
Python
156 lines
4.6 KiB
Python
import uuid
|
|
from collections import deque
|
|
import ljson
|
|
from Crypto.Cipher import AES
|
|
|
|
from .files.tarfile import RewriteableTarFile
|
|
from .files.bunkeredfile import BunkeredFile
|
|
from .backends import components
|
|
|
|
class Bunker(object):
|
|
def __init__(self, path):
|
|
self._tarfile = RewriteableTarFile.open(path)
|
|
self._components = ljson.Table.from_file(self._tarfile.get_file("__bunker_main__"))
|
|
self._stack_of_open_components = deque()
|
|
self._stack_of_open_components.appendleft(deque())
|
|
|
|
@classmethod
|
|
def open(cls, path):
|
|
tarfile = RewriteableTarFile.open(path)
|
|
if(not tarfile.has_file("__bunker_main__")):
|
|
main_file = BunkeredFile.empty("__bunker_main__")
|
|
|
|
header = ljson.Header({"component": {"type": "str", "modifiers": []}
|
|
, "file_name": {"type": "str", "modifiers": []}
|
|
, "type": {"type": "str", "modifiers": []}
|
|
, "key": {"type": "bytes", "modifiers": []}
|
|
, "initvect": {"type": "bytes", "modifiers": []}
|
|
, "length": {"type": "int", "modifiers": []}})
|
|
table = ljson.Table(header, [])
|
|
table.save(main_file)
|
|
tarfile.add_file(main_file)
|
|
tarfile.close()
|
|
|
|
return cls(path)
|
|
|
|
def _load_component(self, name, password):
|
|
chunk_size = 16
|
|
if(not {"component": name} in self._components):
|
|
raise ValueError("unknown component: {}".format(name))
|
|
|
|
info = self._components[{"component": name}]
|
|
file_name = info["file_name"][0]
|
|
length = info["length"][0]
|
|
|
|
password = AES.new(password, AES.MODE_ECB)
|
|
key = password.decrypt(info["key"][0])
|
|
initvect = password.decrypt(info["initvect"][0])
|
|
key = AES.new(key, AES.MODE_CBC, initvect)
|
|
|
|
file_ = self._tarfile.get_file(file_name)
|
|
|
|
decrypted = BunkeredFile.empty(file_name, length_hint=len(file_))
|
|
|
|
length_read = 0
|
|
while(True):
|
|
chunk = file_.read(chunk_size)
|
|
length_read += len(chunk)
|
|
if(not chunk):
|
|
break
|
|
if(length_read == length):
|
|
decrypted.write(key.decrypt(chunk))
|
|
break
|
|
if(length_read > length):
|
|
decrypted.write(key.decrypt(chunk)[:length - length_read])
|
|
break
|
|
decrypted.write(key.decrypt(chunk))
|
|
decrypted.seek(0, 0)
|
|
|
|
return decrypted, info["type"][0]
|
|
|
|
def _add_component(self, name, password, type_):
|
|
if({"name": name} in self._components):
|
|
raise ValueError("component {} is already in the bunker".format(name))
|
|
|
|
|
|
password = AES.new(password, AES.MODE_ECB)
|
|
filename = uuid.uuid1().hex
|
|
key = uuid.uuid4().bytes + uuid.uuid4().bytes
|
|
initvect = uuid.uuid4().bytes
|
|
|
|
key = password.encrypt(key)
|
|
initvect = password.encrypt(initvect)
|
|
|
|
self._components.additem({"component": name
|
|
, "file_name": filename
|
|
, "key": key
|
|
, "initvect": initvect
|
|
, "type": type_
|
|
, "length": 0})
|
|
|
|
self._tarfile.add_file(BunkeredFile.empty(filename))
|
|
self._components.save(self._tarfile.get_file("__bunker_main__"))
|
|
self._tarfile.writeback_file("__bunker_main__")
|
|
|
|
def _save_component(self, name, password, file_):
|
|
chunk_size = 16
|
|
if(not {"component": name} in self._components):
|
|
raise ValueError("unknown component: {}".format(name))
|
|
|
|
info = self._components[{"component": name}]
|
|
file_name = info["file_name"][0]
|
|
|
|
password = AES.new(password, AES.MODE_ECB)
|
|
key = password.decrypt(info["key"][0])
|
|
initvect = password.decrypt(info["initvect"][0])
|
|
key = AES.new(key, AES.MODE_CBC, initvect)
|
|
|
|
file_out = self._tarfile.get_file(file_name)
|
|
file_out.truncate(0)
|
|
file_out.seek(0, 0)
|
|
|
|
file_.seek(0, 0)
|
|
chunk = file_.read(chunk_size)
|
|
length = 0
|
|
while(chunk):
|
|
length += len(chunk)
|
|
if(len(chunk) < 16):
|
|
chunk = chunk + b"\x00"*(16 - len(chunk))
|
|
file_out.write(key.encrypt(chunk))
|
|
chunk = file_.read(chunk_size)
|
|
self._components[{"component": name}]["length"] = length
|
|
|
|
file_out.close()
|
|
|
|
|
|
def __enter__(self):
|
|
self._stack_of_open_components.appendleft(deque())
|
|
return self
|
|
|
|
def __exit__(self):
|
|
open_components = self._stack_of_open_components.popleft()
|
|
for component in open_components:
|
|
component.close(notify_bunker=False)
|
|
return False
|
|
|
|
def add_component(self, type_, name, password):
|
|
if(not type_ in components):
|
|
raise ValueError("unknown component type: '{}'".format(type_))
|
|
component = components[type_].new(self, name, password)
|
|
self._stack_of_open_components[0].append(component)
|
|
return component
|
|
|
|
def get_component(self, name, password):
|
|
type_ = self._components[{"component": name}]["type"][0]
|
|
component = components[type_](self, name, password)
|
|
self._stack_of_open_components[0].append(component)
|
|
|
|
return component
|
|
|
|
def notify_component_is_closed(self, component):
|
|
for context in self._stack_of_open_components:
|
|
for open_component in context:
|
|
if(open_component == component):
|
|
context.remove(open_component)
|
|
break
|