safer import handling

This commit is contained in:
Nick Sweeting 2024-10-08 00:51:58 -07:00
parent 58e37a5e19
commit 46c0463539
No known key found for this signature in database
2 changed files with 34 additions and 6 deletions

View file

@ -8,9 +8,14 @@ from contextlib import contextmanager
DATA_DIR = Path(os.getcwd())
DATA_DIR_STAT = Path(DATA_DIR).stat()
try:
DATA_DIR_STAT = DATA_DIR.stat()
DATA_DIR_UID = DATA_DIR_STAT.st_uid
DATA_DIR_GID = DATA_DIR_STAT.st_gid
except PermissionError:
DATA_DIR_UID = 0
DATA_DIR_GID = 0
DEFAULT_PUID = 911
DEFAULT_PGID = 911
RUNNING_AS_UID = os.getuid()
@ -28,6 +33,29 @@ os.environ.setdefault('PGID', str(DATA_DIR_GID or RUNNING_AS_GID or DEFAULT_PGID
ARCHIVEBOX_USER = int(os.environ['PUID'])
ARCHIVEBOX_GROUP = int(os.environ['PGID'])
if not USER:
try:
# alternative method 1 to get username
import pwd
USER = pwd.getpwuid(ARCHIVEBOX_USER).pw_name
except Exception:
pass
if not USER:
try:
# alternative method 2 to get username
import getpass
USER = getpass.getuser()
except Exception:
pass
if not USER:
try:
# alternative method 3 to get username
USER = os.getlogin() or 'archivebox'
except Exception:
USER = 'archivebox'
#############################################################################################
def drop_privileges():
@ -61,10 +89,9 @@ def SudoPermission(uid=0, fallback=False):
yield
finally:
# then set effective UID back to DATA_DIR owner
DATA_DIR_OWNER = DATA_DIR.stat().st_uid
try:
os.seteuid(DATA_DIR_OWNER)
os.seteuid(ARCHIVEBOX_USER)
except PermissionError as err:
if not fallback:
raise PermissionError(f'Failed to revert uid={uid} back to {DATA_DIR_OWNER} after running code with sudo') from err
raise PermissionError(f'Failed to revert uid={uid} back to {ARCHIVEBOX_USER} after running code with sudo') from err

View file

@ -0,0 +1 @@
__package__ = 'archivebox.misc'