easywks/easywks/files.py
2021-10-03 03:16:31 +02:00

110 lines
3.6 KiB
Python

import os
import fcntl
import stat
from datetime import datetime, timedelta
from pgpy import PGPKey
from .config import Config
from .crypto import create_pgp_key, privkey_to_pubkey
from .util import hash_user_id
def _locked_read(file: str, binary: bool = False):
with open(file, 'r' + 'b' * binary) as f:
fcntl.lockf(f, fcntl.LOCK_SH)
content = f.read()
fcntl.lockf(f, fcntl.LOCK_UN)
return content
def _locked_write(file: str, content, binary: bool = False):
with open(file, 'a' + 'b' * binary) as f:
fcntl.lockf(f, fcntl.LOCK_EX)
f.seek(0)
f.truncate()
f.write(content)
fcntl.lockf(f, fcntl.LOCK_UN)
return content
def make_submission_address_file(domain: str):
return Config[domain].submission_address + '\n'
def make_policy_file(domain: str):
content = f'submission-address: {Config[domain].submission_address}\n'
for flag, value in Config[domain].policy_flags.items():
if not value or len(value) == 0:
content += f'{flag}: {value}\n'
else:
content += flag + '\n'
return content
def init_working_directory():
wdir = Config.working_directory
os.makedirs(wdir, exist_ok=True)
os.chmod(wdir, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
for domain in Config.domains:
# Create necessary files and directories
os.makedirs(os.path.join(wdir, domain, 'hu'), exist_ok=True)
os.makedirs(os.path.join(wdir, domain, 'pending'), exist_ok=True)
_locked_write(os.path.join(wdir, domain, 'submission-address'), make_submission_address_file(domain))
_locked_write(os.path.join(wdir, domain, 'policy'), make_policy_file(domain))
# Create PGP key if it doesn't exist yet
create_pgp_key(domain)
# Export submission key to hu dir
key = privkey_to_pubkey(domain)
uid = hash_user_id(Config[domain].submission_address)
_locked_write(os.path.join(wdir, domain, 'hu', uid), bytes(key), binary=True)
def read_public_key(domain, user):
hu = hash_user_id(user)
keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
key, _ = PGPKey.from_blob(_locked_read(keyfile, binary=True))
return key
def read_hashed_public_key(domain, hu):
keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
key, _ = PGPKey.from_blob(_locked_read(keyfile, binary=True))
return key
def write_public_key(domain, user, key):
hu = hash_user_id(user)
keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
_locked_write(keyfile, bytes(key), binary=True)
def read_pending_key(domain, nonce):
keyfile = os.path.join(Config.working_directory, domain, 'pending', nonce)
key, _ = PGPKey.from_blob(_locked_read(keyfile))
return key
def write_pending_key(domain, nonce, key):
keyfile = os.path.join(Config.working_directory, domain, 'pending', nonce)
_locked_write(keyfile, str(key))
def remove_pending_key(domain, nonce):
keyfile = os.path.join(Config.working_directory, domain, 'pending', nonce)
os.unlink(keyfile)
def clean_stale_requests():
stale = (datetime.utcnow() - timedelta(seconds=Config.pending_lifetime)).timestamp()
for domain in Config.domains:
pending_dir = os.path.join(Config.working_directory, domain, 'pending')
for file in os.listdir(pending_dir):
try:
absfile = os.path.join(pending_dir, file)
if os.stat(absfile).st_mtime < stale:
os.unlink(absfile)
except BaseException as e:
print(e)
continue