Add --upload, --qrcodeHEADmaster

6 files changed, 262 insertions, 54 deletions

diff --git a/api.py b/api.py
new file mode 100644
index 0000000..776395e
--- /dev/null
+++ b/api.py
@@ -0,0 +1,57 @@
+import hmac
+import hashlib
+import traceback
+import itertools
+import time
+import os
+
+from tqdm import tqdm
+from filecrypt import key_id, token_b64encode
+
+_retry_range = lambda retries: itertools.cycle([None]) if retries is None else range(retries)
+
+def upload(path, file_id, size, base_url, api_key, chunk_size=int(10e6), progress=True, max_retries=None):
+    import requests
+
+    with tqdm(total=size, unit='B', unit_scale=True, disable=(not progress)) as pbar, open(path, 'rb') as f:
+        pos = 0
+        while True:
+            chunk = f.read(chunk_size)
+            if not chunk:
+                break
+
+            hash = hashlib.sha3_256()
+            hash.update(chunk)
+            hash = hash.digest()
+
+            content_range = f'bytes {pos}-{pos+len(chunk)-1}/{size}'
+            pos += len(chunk)
+
+            mac = hmac.new(api_key, digestmod='sha3_256')
+            mac.update(file_id.encode())
+            mac.update(hash)
+            mac.update(content_range.encode())
+            
+            post_url = '/'.join([
+                    base_url,
+                    key_id(api_key),
+                    file_id,
+                    token_b64encode(mac.digest()),
+                    token_b64encode(hash)])
+
+            for attempt in _retry_range(max_retries):
+                try:
+                    res = requests.post(post_url, files={'chunk': chunk}, headers={'content-range': content_range})
+                    break
+                except requests.exceptions.RequestException as e:
+                    print(f'Upload error: {type(e).__name__}')
+                    time.sleep(1)
+                    print('Retrying.')
+            else:
+                print(f'Repeated errors uploading. Exiting.')
+                print(f'Leaving encrypted file under: {path}')
+                sys.exit(1)
+            pbar.update(len(chunk))
+
+    os.remove(path)
+
diff --git a/encrypt.py b/encrypt.py
index 6998708..337b3bd 100644
--- a/encrypt.py
+++ b/encrypt.py
@@ -1,50 +1,94 @@
 #!/usr/bin/env python3
 
-if __name__ == '__main__':
-    import os
-    import configparser
-    import argparse
-
-    from tqdm import tqdm
-
-    from filecrypt import generate_keys, payload_size
-
-    parser = argparse.ArgumentParser(description='Filecrypt secure file download encryption tool.'
-            'Encrypts a file for use with the filecrypt server, and output the generated download link.')
-    parser.add_argument('infile')
-    parser.add_argument('-c', '--config', default=None, help='Config file location (default; $XDG_CONFIG_HOME/filecrypt.conf)')
-    parser.add_argument('-b', '--base-url', default=None, help='Base URL for link (also as config option)')
-    parser.add_argument('-f', '--filename', default=None, help='Download filename (default: Same as input filename)')
-    parser.add_argument('-q', '--no-progress', action='store_true', help='Hide progress bar')
-    parser.add_argument('-p', '--progress', action='store_true', help='Show progress bar (default, also as config option)')
-    args = parser.parse_args()
-
-    progress = (not args.no_progress) or args.progress
-    config_path = args.config or os.environ.get('XDG_CONFIG_HOME', os.environ.get('HOME') + '/.config') + '/filecrypt.conf'
-    base_url = args.base_url
-    if os.path.isfile(config_path):
-        with open(config_path) as f:
-            config = configparser.ConfigParser(defaults={'url_base': ''})
-            config.read_string('[DEFAULT]\n'+f.read()) # doesn't parse simple key=value file by default m(
-
-            if base_url is None:
-                base_url = config.get('DEFAULT', 'base_url', fallback='').rstrip('/')
-            if not (args.no_progress or args.progress):
-                progress = config.getboolean('DEFAULT', 'progress', fallback=True)
-
-    if not os.path.isfile(args.infile):
-        print(f'{infile} is not a file or directory, exiting.')
-        os.exit(2)
-
-    download_filename = args.filename or os.path.basename(args.infile)
-
-    file_id, token, encrypt = generate_keys(download_filename)
-    print(f'{base_url}/{file_id}/{token}/{download_filename}')
-
-    if progress:
-        with tqdm(total=payload_size(args.infile), unit='B', unit_scale=True) as pbar:
-            for progress in encrypt(args.infile):
-                pbar.update(progress)
-    else:
-        *encrypt(args.infile),
-    
+if __name__ != '__main__':
+    raise ImportError('Command-line script cannot be imported as module')
+
+import os
+import configparser
+import argparse
+import sys
+
+from tqdm import tqdm
+
+from filecrypt import generate_keys, output_size
+from api import upload
+
+parser = argparse.ArgumentParser(description='Filecrypt secure file download encryption tool.'
+        'Encrypts a file for use with the filecrypt server, and output the generated download link.')
+parser.add_argument('infile')
+parser.add_argument('-c', '--config', default=None, help='Config file location (default; $XDG_CONFIG_HOME/filecrypt.conf)')
+parser.add_argument('-b', '--base-url', default=None, help='Base URL for link (also as config option)')
+parser.add_argument('-f', '--filename', default=None, help='Download filename (default: Same as input filename)')
+parser.add_argument('-n', '--no-progress', action='store_true', help='Hide progress bar')
+parser.add_argument('-p', '--progress', action='store_true', help='Show progress bar (default, also as config option)')
+parser.add_argument('-u', '--upload', action='store_true', help='Upload via HTTP API')
+parser.add_argument('-a', '--api-key', default=None, help='HTTP upload API key')
+parser.add_argument('-q', '--qrcode', action='store_true', help='Show download URL as QR Code')
+parser.add_argument('--upload-chunk-size', type=int, default=None, help='HTTP upload API transfer chunk size')
+parser.add_argument('--max-retries', type=int, default=None, help='HTTP upload request max retries')
+args = parser.parse_args()
+
+progress = (not args.no_progress) or args.progress
+config_path = args.config or os.environ.get('XDG_CONFIG_HOME', os.environ.get('HOME') + '/.config') + '/filecrypt.conf'
+base_url = args.base_url
+api_key = args.api_key
+out_file_size = output_size(args.infile)
+upload_chunk_size = args.upload_chunk_size
+max_retries = args.max_retries
+if os.path.isfile(config_path):
+    with open(config_path) as f:
+        config = configparser.ConfigParser(defaults={'url_base': ''})
+        config.read_string('[DEFAULT]\n'+f.read()) # doesn't parse simple key=value file by default m(
+
+        if base_url is None:
+            base_url = config.get('DEFAULT', 'base_url', fallback='').rstrip('/')
+        if api_key is None:
+            api_key = config.get('DEFAULT', 'api_key', fallback=None)
+        if upload_chunk_size is None:
+            upload_chunk_size = config.get('DEFAULT', 'upload_chunk_size', fallback=None)
+        if max_retries is None:
+           max_retries = config.get('DEFAULT', 'max_retries', fallback=None)
+        if not (args.no_progress or args.progress):
+            progress = config.getboolean('DEFAULT', 'progress', fallback=True)
+
+if not os.path.isfile(args.infile):
+    print(f'{infile} is not a file or directory, exiting.')
+    os.exit(2)
+
+if args.upload:
+    if api_key is None:
+        print(f'HTTP upload API key is required for --upload')
+        ox.exit(2)
+    api_key = api_key.encode()
+
+if upload_chunk_size is None:
+    upload_chunk_size = int(10e6)
+
+download_filename = args.filename or os.path.basename(args.infile)
+
+file_id, token, encrypt = generate_keys(download_filename)
+url = f'{base_url}/{file_id}/{token}/{download_filename}'
+print(url)
+
+if args.qrcode:
+    import qrcode
+    qr = qrcode.QRCode()
+    qr.add_data(url)
+    qr.print_ascii(tty=True)
+
+print('Encrypting...')
+with tqdm(total=out_file_size, unit='B', unit_scale=True, disable=(not progress)) as pbar:
+    for chunk in encrypt(args.infile):
+        pbar.update(len(chunk))
+
+print('Uploading...')
+if args.upload:
+    upload(path = f'{file_id}.enc',
+            file_id=file_id,
+            size=out_file_size,
+            base_url=base_url,
+            chunk_size=upload_chunk_size,
+            progress=progress,
+            api_key=api_key,
+            max_retries=max_retries)
+
diff --git a/filecrypt.py b/filecrypt.py
index 755e857..4b478c6 100644
--- a/filecrypt.py
+++ b/filecrypt.py
@@ -4,11 +4,18 @@ import base64
 import struct
 import subprocess
 import binascii
+import hashlib
 import os
 from contextlib import contextmanager
 
-FILE_ID_LENGTH = 22
-TOKEN_LENGTH = 22
+token_b64decode = lambda token: base64.b64decode(token + '='*((3-len(token)%3)%3), b'+-')
+token_b64encode = lambda token: base64.b64encode(token, b'+-').rstrip(b'=').decode()
+token_b64len    = lambda nbytes: len(token_b64encode(b'0' * nbytes))
+
+key_id = lambda key: token_b64encode(hashlib.sha3_256(b'FILECRYPT_KEY_ID'+key).digest())
+
+FILE_ID_LENGTH = token_b64len(16)
+TOKEN_LENGTH = token_b64len(16)
 HEADER_LENGTH = 56
 
 def generate_keys(download_filename, chunk_size=1000000//16):
@@ -20,7 +27,7 @@ def generate_keys(download_filename, chunk_size=1000000//16):
     token_cipher = AES.new(auth_secret, AES.MODE_GCM)
     token_cipher.update(download_filename.encode())
     ciphertext, token_tag = token_cipher.encrypt_and_digest(key)
-    token = base64.b64encode(ciphertext, b'+-').rstrip(b'=').decode()
+    token = token_b64encode(ciphertext)
 
     def encrypt(filename_in):
         with open(f'{file_id}.enc', 'wb') as fout, open(filename_in, 'rb') as fin:
@@ -37,7 +44,7 @@ def generate_keys(download_filename, chunk_size=1000000//16):
             while block:
                 data = cipher.encrypt(block)
                 fout.write(data)
-                yield len(data)
+                yield data
                 block = fin.read(cipher.block_size*chunk_size)
 
     return file_id, token, encrypt
@@ -45,6 +52,9 @@ def generate_keys(download_filename, chunk_size=1000000//16):
 def payload_size(path):
     return os.stat(path).st_size - HEADER_LENGTH
 
+def output_size(path):
+    return os.stat(path).st_size + HEADER_LENGTH
+
 def decrypt_generator(filename, download_filename, token, seek=0, end=None, chunk_size=1000000//16):
     with open(filename, 'rb') as fin:
         token_nonce = fin.read(16)
@@ -53,7 +63,7 @@ def decrypt_generator(filename, download_filename, token, seek=0, end=None, chun
         data_nonce = fin.read(8)
         assert fin.tell() == HEADER_LENGTH
 
-    ciphertext = base64.b64decode(token + '='*((3-len(token)%3)%3), b'+-')
+    ciphertext = token_b64decode(token)
     token_cipher = AES.new(auth_secret, AES.MODE_GCM, nonce=token_nonce)
     token_cipher.update(download_filename.encode())
     key = token_cipher.decrypt_and_verify(ciphertext, token_tag)
diff --git a/filecrypt_test.conf b/filecrypt_test.conf
new file mode 100644
index 0000000..e9290f6
--- /dev/null
+++ b/filecrypt_test.conf
@@ -0,0 +1,2 @@
+base_url=http://127.0.0.1:5000/
+api_key=foobar
diff --git a/server.py b/server.py
index e729761..ac30f69 100644
--- a/server.py
+++ b/server.py
@@ -3,15 +3,28 @@
 import re
 import os
 
+import hashlib
+import hmac
+
 from flask import Flask, abort, request, Response
 
 import filecrypt
 
 app = Flask(__name__)
+
+# default values
+app.config['MAX_UPLOAD_SIZE'] = int(100e6)
+app.config['MAX_UPLOAD_CHUNK_SIZE'] = int(100e6)
+
 app.config.from_envvar('SECURE_DOWNLOAD_SETTINGS')
+
+upload_keys = {}
+for v in app.config.get('UPLOAD_KEYS', []):
+    upload_keys[filecrypt.key_id(v)] = v
+
 BASE64_RE = re.compile('^[A-Za-z0-9+-_]+=*$')
 
-@app.route('/<file_id>/<token>/<filename>')
+@app.route('/<file_id>/<token>/<filename>', methods=['GET'])
 def download(file_id, token, filename):
     if not BASE64_RE.match(file_id) or len(file_id) != filecrypt.FILE_ID_LENGTH:
         abort(400, 'Invalid file ID format')
@@ -53,3 +66,81 @@ def download(file_id, token, filename):
     response.headers['Content-Disposition'] = f'attachment {filename}'
     return response
 
+@app.route('/<key_id>/<file_id>/<token>/<filehash>', methods=['POST'])
+def upload(key_id, file_id, token, filehash):
+    if not BASE64_RE.match(file_id) or len(file_id) != filecrypt.FILE_ID_LENGTH:
+        abort(400, 'Invalid file ID format')
+    if not BASE64_RE.match(token) or len(token) != filecrypt.token_b64len(32):
+        abort(400, 'Invalid token format')
+    if not BASE64_RE.match(filehash) or len(filehash) != filecrypt.token_b64len(32):
+        abort(400, 'Invalid hash format')
+    if not BASE64_RE.match(key_id) or len(key_id) != filecrypt.token_b64len(32):
+        abort(400, 'Invalid key id format')
+
+    if request.content_length is None:
+        abort(411)
+
+    if request.content_length > app.config['MAX_UPLOAD_CHUNK_SIZE']:
+        abort(413)
+
+    if not key_id in upload_keys:
+        abort(403)
+
+    filehash = filecrypt.token_b64decode(filehash)
+    token = filecrypt.token_b64decode(token)
+    content_range = request.headers.get('Content-Range', 'NO CONTENT RANGE')
+
+    mac = hmac.new(upload_keys[key_id], digestmod='sha3_256')
+    mac.update(file_id.encode())
+    mac.update(filehash)
+    mac.update(content_range.encode())
+    if not hmac.compare_digest(mac.digest(), token):
+        abort(403)
+
+    path = f'{app.config["SERVE_PATH"]}/{file_id}.enc'
+    if os.path.isfile(path):
+        abort(409)
+
+    if 'chunk' not in request.files:
+        abort(400, 'Invalid file payload')
+    data = request.files['chunk'].read()
+
+    hash = hashlib.sha3_256()
+    hash.update(data)
+    if not hmac.compare_digest(hash.digest(), filehash):
+        abort(400)
+
+    tmp_path = f'{path}.uploading'
+    range_header = re.match('^bytes ([0-9]+)-([0-9]+)/([0-9]+|\*)$', content_range)
+    if not range_header:
+        if os.path.isfile(tmp_path):
+            os.remove(tmp_path)
+
+        with open(path, 'wb') as f:
+            f.write(data)
+        print(f'{request.remote_addr}: {file_id} UPLOAD')
+        return 'success', 200
+
+    else:
+        range_start, range_end, size = range_header.groups()
+        if size == '*':
+            abort(400, 'Content-range header if used must include total size')
+        try:
+            range_start, range_end, size = int(range_start), int(range_end), int(size)
+        except ValueError:
+            abort(400)
+
+        with open(tmp_path, 'ab') as f:
+            if range_start > f.tell():
+                abort(416)
+
+            f.truncate(range_start)
+            f.write(data)
+
+        if range_end+1 == size:
+            os.rename(tmp_path, path)
+            print(f'{request.remote_addr}: {file_id} UPLOAD')
+            return 'success', 200
+        
+        return 'partial', 206
+
diff --git a/test_config.cfg b/test_config.cfg
new file mode 100644
index 0000000..2ac0560
--- /dev/null
+++ b/test_config.cfg
@@ -0,0 +1,4 @@
+SERVE_PATH = '/tmp'
+UPLOAD_KEYS = [
+    b'foobar' # test key
+]