diff options
author | jerry73204 <jerry73204@gmail.com> | 2019-06-10 23:18:00 +0800 |
---|---|---|
committer | jerry73204 <jerry73204@gmail.com> | 2019-06-10 23:18:00 +0800 |
commit | 4f816b8f9eeafcf7974aeaf6f052e68b6f351507 (patch) | |
tree | 82c9752d9f80c46359c7ac0eddceba94f5e4b657 | |
parent | 67cb88e2a47a7b46567540a2498a210cb82db7a4 (diff) | |
download | cns-final-tor-store-4f816b8f9eeafcf7974aeaf6f052e68b6f351507.tar.gz cns-final-tor-store-4f816b8f9eeafcf7974aeaf6f052e68b6f351507.tar.zst cns-final-tor-store-4f816b8f9eeafcf7974aeaf6f052e68b6f351507.zip |
Done virtual file system and simple shell
-rw-r--r-- | demo/demo.jpg | bin | 0 -> 962 bytes | |||
-rw-r--r-- | demo/random_block | 1 | ||||
-rw-r--r-- | demo/zero_block | bin | 0 -> 100 bytes | |||
-rw-r--r-- | fs.py | 143 | ||||
-rwxr-xr-x | main.py | 13 | ||||
-rw-r--r-- | src/fs.py | 505 | ||||
-rwxr-xr-x | src/main.py | 280 | ||||
-rwxr-xr-x | src/tor_helper.py | 58 | ||||
-rw-r--r-- | src/tor_utils.py (renamed from tor_utils.py) | 45 |
9 files changed, 861 insertions, 184 deletions
diff --git a/demo/demo.jpg b/demo/demo.jpg Binary files differnew file mode 100644 index 0000000..626a1d7 --- /dev/null +++ b/demo/demo.jpg diff --git a/demo/random_block b/demo/random_block new file mode 100644 index 0000000..aae1c33 --- /dev/null +++ b/demo/random_block @@ -0,0 +1 @@ +e( w#fɪzK$d$&atڬ毑{":9[jS? 1կQrcL?LuMj<{l\g uh
\ No newline at end of file diff --git a/demo/zero_block b/demo/zero_block Binary files differnew file mode 100644 index 0000000..eeb5760 --- /dev/null +++ b/demo/zero_block @@ -1,143 +0,0 @@ -import re -import tor_utils - - -class VFS: - def __init__(self, data_size=800, replica_factor=3): - assert data_size / 8 == data_size // 8 - self.fs = dict() - self.re_slash = re.compile('/+') - self.block_size = data_size // 8 - self.replica_factor = replica_factor - - def open(self, path): - path = self.re_slash.sub('/', path) - tokens = path.split('/') - handle = self.fs - - for tk in tokens[:-1]: - if not tk: - raise ValueError('Invalid path %s' % path) - handle = handle.get(tk, dict()) - if not isinstance(handle, dict): - raise ValueError('%s is not valid' % path) - - default_file_handle = FileHandle( - block_size=self.block_size, - replica_factor=self.replica_factor, - ) - file_handle = handle.get(tokens[-1], default_file_handle) - return file_handle - - -class FileHandle: - def __init__(self, block_size, replica_factor): - self.block_size = block_size - self.replica_factor = replica_factor - self.block_store = dict() - - def load_block(self, index, check_boundary=True): - assert index >= 0 - - if index in self.block_store: - # Load from one of its replica - for addr in self.block_store[index]: - block = tor_utils.load_block(addr) - if block is not None: - return block - - # Fail - raise ValueError("Fail to load block at index %d" % index) - - elif check_boundary: - raise ValueError("Index out of bound") - - else: - return b'\x00' * self.block_size - - def store_block(self, index, block): - assert index >= 0 - - addresses = list() - for _ in range(self.replica_factor): - addr = tor_utils.store_block(block, self.block_size) - addresses.append(addr) - - self.block_store[index] = addresses - - def read(self, offset, length): - begin_offset = offset - end_offset = offset + length - - begin_index = begin_offset // self.block_size - end_index = end_offset // self.block_size - - # Load first block - if begin_offset / self.block_size != begin_index: - block = self.load_block(begin_index) - front_length = self.block_size * (begin_index + 1) - begin_offset - assert front_length > 0 - front_block = block[:-front_length] - begin_index += 1 - else: - front_length = 0 - front_block = b'' - - # Load last block - if end_offset / self.block_size != end_index: - block = self.load_block(end_index) - tail_length = self.block_size * (end_index) - end_offset - assert tail_length > 0 - tail_block = block[:tail_length] - else: - tail_length = 0 - tail_block = b'' - - # Load intermediate blocks - data = front_block - - for index in range(begin_index, end_index): - block = self.load_block(index) - data += block - - data += tail_block - return data - - def write(self, offset, data): - length = len(data) - begin_offset = offset - end_offset = offset + length - - begin_index = begin_offset // self.block_size - end_index = end_offset // self.block_size - - # Update first block - if begin_offset / self.block_size != begin_index: - block = self.load_block(begin_index) - front_length = self.block_size * (begin_index + 1) - begin_offset - assert front_length > 0 - - new_block = block[:-front_length] + data[:front_length] - self.store_block(begin_index, new_block) - begin_index += 1 - else: - front_length = 0 - - # Update last block - if end_offset / self.block_size != end_index: - block = self.load_block(end_index) - tail_length = self.block_size * (end_index) - end_offset - assert tail_length > 0 - - new_block = data[-tail_length:] + block[:tail_length] - self.store_block(end_index, new_block) - else: - tail_length = 0 - - # Update intermediate blocks - for index in range(begin_index, end_index): - begin_data_offset = front_length + self.block_size * (index - begin_index) - end_data_offset = begin_data_offset + self.block_size - - new_block = data[begin_data_offset:end_data_offset] - self.store_block(index, new_block) diff --git a/main.py b/main.py deleted file mode 100755 index 738c9db..0000000 --- a/main.py +++ /dev/null @@ -1,13 +0,0 @@ -#!/usr/bin/env python3 -import argparse - - -def main(): - parser = argparse.ArgumentParser() - args = parser.parse_args() - - - - -if __name__ == '__main__': - main() diff --git a/src/fs.py b/src/fs.py new file mode 100644 index 0000000..717ebe8 --- /dev/null +++ b/src/fs.py @@ -0,0 +1,505 @@ +import os +import re +import shutil +from json import JSONEncoder, JSONDecoder +import multiprocessing as mp +import asyncio +import sys + +from logzero import logger + + +class VFS: + def __init__(self, data_size=800, replica_factor=1, max_workers=None, buf_size=2**30): + assert data_size / 8 == data_size // 8 + self.re_slash = re.compile('/+') + self.block_length = data_size // 8 + self.replica_factor = replica_factor + self.buf_size = buf_size + self.fs = dict() + + if max_workers is None: + max_workers = mp.cpu_count() * 2 + + self.queue = asyncio.Queue(maxsize=max_workers) + + def parse_path(self, path): + path = self.re_slash.sub('/', path) + tokens = path.split('/') + + normalized_tokens = list() + + for tk in tokens: + if not tk: + raise VFSError('Invalid path %s' % path) + elif tk == '.': + continue + elif tk == '..': + if normalized_tokens: + normalized_tokens.pop() + else: + normalized_tokens.append(tk) + + return normalized_tokens + + def traverse(self, path, tokens): + parent_handle = self.fs + handle = self.fs + + for tk in tokens: + if isinstance(handle, FileHandle): + raise VFSError('%s is not a file or directory' % path) + elif tk not in handle: + raise VFSError('%s is not a file or directory' % path) + else: + parent_handle = handle + handle = handle[tk] + + return parent_handle, handle + + def list(self, path): + tokens = self.parse_path(path) + parent_handle, handle = self.traverse(path, tokens) + + if isinstance(handle, FileHandle): + return [path] + else: + assert isinstance(handle, dict) + return list(handle.keys()) + + def find(self, path): + tokens = self.parse_path(path) + parent_handle, handle = self.traverse(path, tokens) + + def recursive_list(prefix, h): + if isinstance(h, FileHandle): + return [prefix] + else: + assert isinstance(h, dict) + ret = [prefix] + for name, child in h.items(): + child_path = '%s/%s' % (prefix, name) + ret += recursive_list(child_path, child) + return ret + + result = recursive_list(path, handle) + return result + + def touch(self, path): + tokens = self.parse_path(path) + parent_handle, handle = self.traverse(path, tokens[:-1]) + + if isinstance(handle, FileHandle): + raise VFSError('%s is not valid' % path) + elif tokens[-1] not in handle: + handle[tokens[-1]] = FileHandle( + block_length=self.block_length, + replica_factor=self.replica_factor, + queue=self.queue, + ) + + def mkdir(self, path): + tokens = self.parse_path(path) + parent_handle, handle = self.traverse(path, tokens[:-1]) + + if isinstance(handle, FileHandle): + raise VFSError('%s is not valid' % path) + elif tokens[-1] in handle: + raise VFSError('%s already exists' % path) + else: + handle[tokens[-1]] = dict() + + def remove(self, path, recursive=False): + tokens = self.parse_path(path) + parent_handle, handle = self.traverse(path, tokens) + + if not recursive and isinstance(handle, dict): + raise VFSError('%s is not a file' % path) + + if tokens: + parent_handle.pop(tokens[-1], None) + else: + self.fs = dict() + + async def copy(self, from_path, to_path): + from_outer = False + if from_path[0] == '@': + from_path = from_path[1:] + from_outer = True + else: + from_tokens = self.parse_path(from_path) + _, from_handle = self.traverse(from_path, from_tokens) + if not isinstance(from_handle, FileHandle): + raise VFSError('Cannot copy from %s' % from_path) + + to_outer = False + if to_path[0] == '@': + to_path = to_path[1:] + to_outer = True + else: + to_tokens = self.parse_path(to_path) + _, to_parent_handle = self.traverse(to_path, to_tokens[:-1]) + if not isinstance(to_parent_handle, dict): + raise VFSError('Cannot copy to %s' % to_path) + + to_handle = FileHandle( + block_length=self.block_length, + replica_factor=self.replica_factor, + queue=self.queue, + ) + to_parent_handle[to_tokens[-1]] = to_handle + + if from_outer: + if to_outer: + shutil.copyfile(from_path, to_path) + else: + with open(from_path, 'rb') as from_file: + offset = 0 + while True: + buf = from_file.read(self.buf_size) + if not buf: + break + await to_handle.write(offset, buf) + offset += len(buf) + + else: + if to_outer: + with open(to_path, 'wb') as to_file: + for offset in range(0, from_handle.file_length, self.buf_size): + buf = await from_handle.read(offset, self.buf_size) + if not buf: + raise VFSError('Unexpected EOF') + to_file.write(buf) + else: + for offset in range(0, from_handle.file_length, self.buf_size): + buf = await from_handle.read(offset, self.buf_size) + await to_handle.write(offset, buf) + + def stat(self, path): + tokens = self.parse_path(path) + _, handle = self.traverse(path, tokens) + + if isinstance(handle, FileHandle): + return { + 'type': 'file', + 'size': handle.file_length, + 'tor_addresses': handle.block_store, + } + else: + assert isinstance(handle, dict) + return { + 'type': 'directory', + } + + def open(self, path): + tokens = self.parse_path(path) + _, handle = self.traverse(path, tokens) + + if not isinstance(handle, FileHandle): + raise VFSError('%s is not a file' % path) + + return handle + + +class FileHandle: + def __init__(self, block_length, queue, replica_factor=1): + self.block_length = block_length + self.replica_factor = replica_factor + self.block_store = dict() + self.file_length = 0 + self.queue = queue + self.tor_helper_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + 'tor_helper.py', + ) + + async def load_block(self, index, check_boundary=True): + logger.debug('load_block(%d, check_consistency=%d)', index, check_boundary) + assert index >= 0 + + if index in self.block_store: + # Load from one of its replica + for addr in self.block_store[index]: + logger.info('Loading replica from Onion address %s.onion', addr) + await self.queue.put(None) # Constraint # of concurrent workers + proc = await asyncio.create_subprocess_exec( + sys.executable, + self.tor_helper_path, + 'load', + '1024', + '800', + '200', + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + ) + proc.stdin.write(bytes('%s\n' % addr, 'ASCII')) + block = await proc.stdout.read(100) + await proc.wait() + await self.queue.get() + + if block: + assert len(block) == self.block_length + return block + else: + logger.warning('Failed to load replica from Onion address %s.onion', addr) + + raise VFSError("Fail to load block at index %d" % index) + + elif check_boundary: + raise VFSError("Index out of bound") + + else: + return b'\x00' * self.block_length + + async def store_block(self, index, block): + logger.debug('store_block(%d, ..)', index) + assert index >= 0 and len(block) <= self.block_length + + # Pad block + if len(block) < self.block_length: + block = block + b'\x00' * (self.block_length - len(block)) + + futures = list() + processes = list() + + for replica_index in range(self.replica_factor): + logger.info('Storing replica %d/%d for block index %d', replica_index + 1, self.replica_factor, index) + + await self.queue.put(None) # Constraint # of concurrent workers + proc = await asyncio.create_subprocess_exec( + sys.executable, + self.tor_helper_path, + 'store', + '1024', + '800', + '200', + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + ) + proc.stdin.write(block) + addr_future = proc.stdout.readline() + futures.append((proc, addr_future)) + + addresses = list() + for proc, addr_future in futures: + addr = str(await addr_future, 'ASCII')[:-1] # Strip '\n' + addresses.append(addr) + await proc.wait() + await self.queue.get() + + self.block_store[index] = addresses + + async def read(self, offset, length): + begin_offset = offset + end_offset = offset + length + + # Sanitize boundary + if begin_offset >= self.file_length: + return b'' + + if end_offset > self.file_length: + end_offset = self.file_length + + begin_index = begin_offset // self.block_length + end_index = end_offset // self.block_length + + has_front = begin_offset / self.block_length != begin_index + has_tail = end_offset / self.block_length != end_index + + # Single block case + if begin_index == end_index: + block = await self.load_block(begin_index) + front_strip = begin_offset - self.block_length * begin_index + tail_strip = self.block_length * (end_index + 1) - end_offset + return block[front_strip:-tail_strip] if tail_strip > 0 \ + else block[front_strip:] + + # Load first block + if has_front: + block = await self.load_block(begin_index) + front_length = self.block_length * (begin_index + 1) - begin_offset + assert front_length > 0 + front_block = block[:-front_length] + begin_index += 1 + else: + front_block = b'' + + # Load last block + if has_tail: + block = await self.load_block(end_index) + tail_length = end_offset - self.block_length * end_index + assert tail_length > 0 + tail_block = block[:tail_length] + else: + tail_block = b'' + + # Load intermediate blocks + data = front_block + + for index in range(begin_index, end_index): + block = await self.load_block(index) + data += block + + data += tail_block + return data + + async def write(self, offset, data): + length = len(data) + begin_offset = offset + end_offset = offset + length + + # Update file size + if end_offset > self.file_length: + self.file_length = end_offset + + begin_index = begin_offset // self.block_length + end_index = end_offset // self.block_length + + has_front = begin_offset / self.block_length != begin_index + has_tail = end_offset / self.block_length != end_index + + # Single block case + if begin_index == end_index: + block = await self.load_block(begin_index, check_boundary=False) + front_strip = begin_offset - begin_index * self.block_length + tail_strip = self.block_length - (front_strip + self.block_length) + front_block = block[:front_strip] + tail_block = block[-tail_strip:] if tail_strip > 0 else b'' + new_block = front_block + data + tail_block + await self.store_block(begin_index, new_block) + return + + # Store blocks asynchrnously + futures = list() + + # Update first block + if has_front: + block = await self.load_block(begin_index, check_boundary=False) + front_length = self.block_length * (begin_index + 1) - begin_offset + assert front_length > 0 + + new_block = block[:-front_length] + data[:front_length] + future = self.store_block(begin_index, new_block) + futures.append(future) + begin_index += 1 + else: + front_length = 0 + + # Update last block + if has_tail: + block = await self.load_block(end_index, check_boundary=False) + tail_length = end_offset - self.block_length * end_index + assert tail_length > 0 + + new_block = data[-tail_length:] + block[tail_length:] + future = self.store_block(end_index, new_block) + futures.append(future) + else: + tail_length = 0 + + # Update intermediate blocks + for index in range(begin_index, end_index): + begin_data_offset = front_length + self.block_length * (index - begin_index) + end_data_offset = begin_data_offset + self.block_length + + new_block = data[begin_data_offset:end_data_offset] + future = self.store_block(index, new_block) + futures.append(future) + + await asyncio.gather(*futures) + + +class VFSError(Exception): + pass + + +class VFSJsonEncoder(JSONEncoder): + def default(self, obj): + if isinstance(obj, VFS): + return self.serialize_vfs(obj) + elif isinstance(obj, FileHandle): + return self.serialize_filehandle(obj) + else: + return super(VFSJsonEncoder, self).default(obj) + + def serialize_vfs(self, obj): + assert isinstance(obj, VFS) + + def recursive_serialize_handle(handle): + if isinstance(handle, FileHandle): + return self.serialize_filehandle(handle) + else: + assert isinstance(handle, dict) + result = dict() + for name, child in handle.items(): + result[name] = recursive_serialize_handle(child) + return result + + fs_obj = recursive_serialize_handle(obj.fs) + result = { + '_type': 'VFS', + 'block_length': obj.block_length, + 'replica_factor': obj.replica_factor, + 'buf_size': obj.buf_size, + 'fs': fs_obj, + } + return result + + def serialize_filehandle(self, obj): + result = { + '_type': 'FileHandle', + 'block_length': obj.block_length, + 'replica_factor': obj.replica_factor, + 'file_length': obj.file_length, + 'block_store': obj.block_store, + } + return result + + +class VFSJsonDecoder(JSONDecoder): + def __init__(self, *args, **kargs): + JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kargs) + + def object_hook(self, obj): + if '_type' not in obj: + return obj + + type_ = obj['_type'] + + if type_ == 'VFS': + vfs = VFS() + vfs.block_length = obj['block_length'] + vfs.replica_factor = obj['replica_factor'] + vfs.buf_size = obj['buf_size'] + + def recursive_deserialize_handle(handle): + if isinstance(handle, FileHandle): + handle.queue = vfs.queue + return handle + else: + assert isinstance(handle, dict) + result = dict() + for name, child in handle.items(): + result[name] = recursive_deserialize_handle(child) + + return result + + vfs.fs = recursive_deserialize_handle(obj['fs']) + return vfs + + elif type_ == 'FileHandle': + return self.decode_filehandle(obj) + + return obj + + def decode_filehandle(self, obj): + dummy_queue = None + handle = FileHandle(obj['block_length'], dummy_queue, replica_factor=obj['replica_factor']) + handle.file_length = obj['file_length'] + + block_store = dict() + for index, replicas in obj['block_store'].items(): + block_store[int(index)] = replicas + + handle.block_store = block_store + return handle diff --git a/src/main.py b/src/main.py new file mode 100755 index 0000000..fa9869a --- /dev/null +++ b/src/main.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +import argparse +import readline +import os +import atexit +import json +import pprint +import asyncio + +import logzero + +import fs + + +async def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--histfile', default='.torfs_history') + parser.add_argument('--vfs', default='.torfs_vfs') + + args = parser.parse_args() + + # Configure logger + if 'LOGLEVEL' in os.environ: + logzero.loglevel(os.environ['LOGLEVEL']) + + # Configure readline + try: + readline.read_history_file(args.histfile) + except FileNotFoundError: + pass + + readline.parse_and_bind('tab: complete') + atexit.register(readline.write_history_file, args.histfile) + + # Load saved file system + if os.path.isfile(args.vfs): + with open(args.vfs, 'r') as file_vfs: + vfs = json.load(file_vfs, cls=fs.VFSJsonDecoder) + else: + vfs = fs.VFS() + + open_files = dict() + fd_count = 0 + + # Serve user commands + while True: + try: + command = input('torfs> ') + except EOFError: + break + except KeyboardInterrupt: + print() + continue + + tokens = command.split() + program = tokens[0] + + try: + if program == 'ls': + if len(tokens) > 2: + print('Invalid arguments') + continue + + if len(tokens) == 2: + path = tokens[1] + else: + path = '.' + + for name in vfs.list(path): + print(name) + + elif program == 'find': + if len(tokens) > 2: + print('Invalid arguments') + continue + + if len(tokens) == 2: + path = tokens[1] + else: + path = '.' + + for name in vfs.find(path): + print(name) + + elif program == 'touch': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + vfs.touch(path) + + elif program == 'mkdir': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + vfs.mkdir(path) + + elif program == 'rm': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + vfs.remove(path) + + elif program == 'rmdir': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + vfs.remove(path, recursive=True) + + elif program == 'cp': + if len(tokens) != 3: + print('Invalid arguments') + continue + + from_path = tokens[1] + to_path = tokens[2] + await vfs.copy(from_path, to_path) + + elif program == 'stat': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + file_stat = vfs.stat(path) + pprint.pprint(file_stat) + + elif program == 'open': + if len(tokens) != 2: + print('Invalid arguments') + continue + + path = tokens[1] + fp = vfs.open(path) + open_files[fd_count] = fp + print('fd = %d' % fd_count) + fd_count += 1 + + elif program == 'fd': + if len(tokens) != 1: + print('"fd" command has no arguments') + continue + + for fd in open_files.keys(): + print(fd) + + elif program == 'close': + if len(tokens) != 2: + print('Invalid arguments') + continue + + try: + fd = int(tokens[1]) + except ValueError: + print('Invalid arguments') + continue + + if fd not in open_files: + print('Invalid arguments') + + open_files.remove(fd) + + elif program == 'read': + if len(tokens) != 4: + print('Invalid arguments') + continue + + try: + fd = int(tokens[1]) + offset = int(tokens[2]) + length = int(tokens[3]) + except ValueError: + print('Invalid arguments') + continue + + if fd not in open_files: + print('Invalid arguments') + continue + + fp = open_files[fd] + buf = await fp.read(offset, length) + print(buf) + + elif program == 'write': + if len(tokens) != 4: + print('Invalid arguments') + continue + + try: + fd = int(tokens[1]) + offset = int(tokens[2]) + data = eval(tokens[3]) + except ValueError: + print('Invalid arguments') + continue + + if isinstance(data, bytes): + print('Invalid arguments') + continue + + if fd not in open_files: + print('Invalid arguments') + continue + + fp = open_files[fd] + buf = await fp.write(offset, data) + + elif program == 'exit': + exit(0) + + elif program == 'help': + if len(tokens) != 1: + print('"help" command has no arguments') + continue + + print(r'''ls [PATH] + List directory. + +find [PATH] + Recursively list files and directories. + +touch PATH + Create empty file. + +mkdir PATH + Create directory. + +rm PATH + Delete file. + +rmdir PATH + Recursively delete directory. + + +cp FROM_PATH TO_PATH + Copy files. + If FROM_PATH or TO_PATH is prefixed with '@', it indicates + the path on host. For example, "cp @from.jpg to.jpg" reads + "from.jpg" from host, and copies to "to.jpg" in TorFS. + +stat PATH + Show file information. + +open PATH + Open a file in TorFS and allocate a file descriptor. + +fd + List open files. + +close FD + Close a file descriptor. + +read FD OFFSET LENGTH + Read data with LENGTH in size starting from OFFSET on FD. + +write FD OFFSET DATA + Write data starting from OFFSET on FD. The DATA is encoded + in Python bytes representation (b'\x00\x01...'). +''') + + else: + print('Invalid command') + + except fs.VFSError as e: + print('Error:', e) + + # Save file system state + with open(args.vfs, 'w') as file_vfs: + json.dump(vfs, file_vfs, cls=fs.VFSJsonEncoder) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/src/tor_helper.py b/src/tor_helper.py new file mode 100755 index 0000000..15e8eaf --- /dev/null +++ b/src/tor_helper.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +import argparse +import sys +import os + +import logzero + +import tor_utils + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('COMMAND', choices=['store', 'load']) + parser.add_argument('KEY_SIZE', type=int) + parser.add_argument('DATA_SIZE', type=int) + parser.add_argument('NONCE_SIZE', type=int) + args = parser.parse_args() + + # Configure logger + if 'LOGLEVEL' in os.environ: + logzero.loglevel(os.environ['LOGLEVEL']) + + # Parse arguments + data_length = args.DATA_SIZE // 8 + assert args.DATA_SIZE / 8 == data_length + + if args.COMMAND == 'load': + addr = input() + block = tor_utils.load_block( + addr, + key_size=args.KEY_SIZE, + data_size=args.DATA_SIZE, + nonce_size=args.NONCE_SIZE, + ) + if block is not None: + sys.stdout.buffer.write(block) + else: + exit(1) + + elif args.COMMAND == 'store': + block = sys.stdin.buffer.read(data_length) + addr = tor_utils.store_block( + block, + key_size=args.KEY_SIZE, + data_size=args.DATA_SIZE, + nonce_size=args.NONCE_SIZE, + ) + if addr is not None: + print(addr) + else: + exit(1) + + else: + raise ValueError('Command %s is not understood' % args.COMMAND) + + +if __name__ == '__main__': + main() diff --git a/tor_utils.py b/src/tor_utils.py index 8af2e23..39bbc0d 100644 --- a/tor_utils.py +++ b/src/tor_utils.py @@ -5,12 +5,16 @@ import asn1 import gmpy2 from Crypto.PublicKey import RSA from stem.control import Controller +from logzero import logger -def store_block(data, data_len=100): +def store_block(data, key_size=1024, data_size=800, nonce_size=200): + data_length = data_size // 8 + assert data_size / 8 == data_length + assert len(data) == data_length + # Generate RSA key pair - data = random.getrandbits(data_len * 8).to_bytes(data_len, 'little') - key = forge_rsa_key(data) + key = forge_rsa_key(data, key_size=key_size, data_size=data_size, nonce_size=nonce_size) # Public hidden service with Controller.from_port() as controller: @@ -22,6 +26,7 @@ def store_block(data, data_len=100): key_content=key, ) + logger.info('Hidden service ID %s published', response.service_id) return response.service_id @@ -73,8 +78,12 @@ def forge_rsa_key(data: bytes, key_size=1024, data_size=800, nonce_size=200): key = RSA.construct((n, e, d), consistency_check=True) except ValueError: continue + break + logger.debug('Created public key with n = %s', bin(n)) + assert data_from_public_key(n) == data + # Tor accepts DER-encoded, then base64 encoded RSA key # https://github.com/torproject/tor/blob/a462ca7cce3699f488b5e2f2921738c944cb29c7/src/feature/control/control_cmd.c#L1968 der = key.export_key('DER', pkcs=1) @@ -82,30 +91,7 @@ def forge_rsa_key(data: bytes, key_size=1024, data_size=800, nonce_size=200): return ret -def load_block(block_id, data_len=61): - url = '%s.union' % block_id - - with Controller.from_port() as controller: - controller.authenticate() - try: - descriptor = str(controller.get_hidden_service_descriptor(url)) - except ValueError: - return None - - print(descriptor) - - # b = descriptor[descriptor.find('BEGIN'):descriptor.find('END')] - # c = descriptor[descriptor.find('BEGIN MESSAGE'):descriptor.find('END MESSAGE')] - #print(x, a) - # print(x, hashlib.sha256((b + c).encode()).hexdigest()) - - -def data_from_public_key(n, key_size=1024, data_size=488): - data_num = (n - (1 << (key_size - 1))) >> ((key_size - 1) - data_size) - return data_num.to_bytes(data_size // 8, 'little') - - -def load_block(name: str): +def load_block(name: str, key_size=1024, data_size=800, nonce_size=200): with Controller.from_port() as controller: controller.authenticate() a = str(controller.get_hidden_service_descriptor(name)) @@ -113,9 +99,12 @@ def load_block(name: str): decoder = asn1.Decoder() decoder.start(base64.b64decode(public)) decoder.start(decoder.read()[1]) - data = data_from_public_key(decoder.read()[1]) + n = decoder.read()[1] + logger.debug('Received public key with n = %s', bin(n)) + data = data_from_public_key(n, key_size=key_size, data_size=data_size, nonce_size=nonce_size) return data + def data_from_public_key(n, key_size=1024, data_size=800, nonce_size=200): data_num = (n - (1 << (key_size - 1))) >> (key_size - 1 - data_size) return data_num.to_bytes(data_size // 8, 'little') |