Skip to content

Harden IMAP hook against symlink-followed attachment writes #67075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Lougarou wants to merge 1 commit into from
+114 −34
Closed

Harden IMAP hook against symlink-followed attachment writes #67075

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion providers/imap/src/airflow/providers/imap/hooks/imap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,22 @@ def _correct_path(self, name: str, local_output_directory: str) -> str:
def _create_file(self, name: str, payload: Any, local_output_directory: str) -> None:
file_path = self._correct_path(name, local_output_directory)

with open(file_path, "wb") as file:
# Defense-in-depth: ``_is_symlink`` (called from ``_create_files``) only inspects
# ``name`` relative to the CWD, not the actual target the attachment is written
# to. Re-check the real joined target and refuse to write through a symlink there,
# matching the existing ``_is_symlink`` rejection style.
if os.path.islink(file_path):
self.log.error("Can not create file because it is a symlink!")
return

# ``O_NOFOLLOW`` closes the TOCTOU window between the check above and the open:
# if ``file_path`` is (or becomes) a symlink, the open fails instead of following
# the link. It is feature-gated with ``getattr`` because it does not exist on
# Windows. ``O_CREAT | O_TRUNC`` (deliberately without ``O_EXCL``) preserves the
# hook's existing overwrite-on-redownload behaviour.
flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | getattr(os, "O_NOFOLLOW", 0)
fd = os.open(file_path, flags, 0o600)
with os.fdopen(fd, "wb") as file:
file.write(payload)


Expand Down
131 changes: 98 additions & 33 deletions providers/imap/tests/unit/imap/hooks/test_imap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import imaplib
import json
import os
from unittest.mock import Mock, mock_open, patch

import pytest
Expand All @@ -30,7 +31,11 @@
from tests_common.test_utils.config import conf_vars

imaplib_string = "airflow.providers.imap.hooks.imap.imaplib"
open_string = "airflow.providers.imap.hooks.imap.open"
# ``ImapHook._create_file`` writes attachments through ``os.open`` / ``os.fdopen``
# (with ``O_NOFOLLOW``) rather than the builtin ``open`` so that a symlink at the
# real target path cannot redirect the write. Tests patch those two symbols.
os_open_string = "airflow.providers.imap.hooks.imap.os.open"
os_fdopen_string = "airflow.providers.imap.hooks.imap.os.fdopen"


def _create_fake_imap(mock_imaplib, with_mail=False, attachment_name="test1.csv", use_ssl=True):
Expand Down Expand Up @@ -305,45 +310,49 @@ def test_retrieve_mail_attachments_with_mail_filter(self, mock_imaplib):

mock_imaplib.IMAP4_SSL.return_value.search.assert_called_once_with(None, mail_filter)

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_found(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_found(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments("test1.csv", "test_directory")

mock_open_method.assert_called_once_with("test_directory/test1.csv", "wb")
mock_open_method.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")
assert mock_os_open.call_args[0][0] == "test_directory/test1.csv"
mock_fdopen.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_not_found(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_not_found(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)

with ImapHook() as imap_hook:
with pytest.raises(AirflowException):
imap_hook.download_mail_attachments("test1.txt", "test_directory")

mock_open_method.assert_not_called()
mock_open_method.return_value.write.assert_not_called()
mock_os_open.assert_not_called()
mock_fdopen.return_value.write.assert_not_called()

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_regex_found(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_with_regex_found(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(
name=r"test(\d+).csv", local_output_directory="test_directory", check_regex=True
)

mock_open_method.assert_called_once_with("test_directory/test1.csv", "wb")
mock_open_method.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")
assert mock_os_open.call_args[0][0] == "test_directory/test1.csv"
mock_fdopen.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_regex_not_found(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_with_regex_not_found(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)

with ImapHook() as imap_hook:
Expand All @@ -354,49 +363,105 @@ def test_download_mail_attachments_with_regex_not_found(self, mock_imaplib, mock
check_regex=True,
)

mock_open_method.assert_not_called()
mock_open_method.return_value.write.assert_not_called()
mock_os_open.assert_not_called()
mock_fdopen.return_value.write.assert_not_called()

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_latest_only(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_with_latest_only(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(
name="test1.csv", local_output_directory="test_directory", latest_only=True
)

mock_open_method.assert_called_once_with("test_directory/test1.csv", "wb")
mock_open_method.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")
assert mock_os_open.call_args[0][0] == "test_directory/test1.csv"
mock_fdopen.return_value.write.assert_called_once_with(b"SWQsTmFtZQoxLEZlbGl4")

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_escaping_chars(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_with_escaping_chars(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True, attachment_name="../test1.csv")

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(name="../test1.csv", local_output_directory="test_directory")

mock_open_method.assert_not_called()
mock_open_method.return_value.write.assert_not_called()
mock_os_open.assert_not_called()
mock_fdopen.return_value.write.assert_not_called()

@patch("airflow.providers.imap.hooks.imap.os.path.islink", return_value=True)
@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_symlink(self, mock_imaplib, mock_open_method, mock_is_symlink):
def test_download_mail_attachments_with_symlink(
self, mock_imaplib, mock_os_open, mock_fdopen, mock_is_symlink
):
_create_fake_imap(mock_imaplib, with_mail=True, attachment_name="symlink")

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(name="symlink", local_output_directory="test_directory")

assert mock_is_symlink.call_count == 1
mock_open_method.assert_not_called()
mock_open_method.return_value.write.assert_not_called()
# ``os.path.islink`` is consulted twice: once by ``_is_symlink`` on the bare
# ``name`` and once by ``_create_file`` on the real joined target path.
assert mock_is_symlink.call_count >= 1
mock_os_open.assert_not_called()
mock_fdopen.return_value.write.assert_not_called()

@patch("airflow.providers.imap.hooks.imap.os.path.islink")
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_symlink_at_real_target(
self, mock_imaplib, mock_os_open, mock_fdopen, mock_islink
):
"""A symlink at the *real* joined output path must block the write.

``_is_symlink`` (called from ``_create_files``) only inspects ``name``
relative to the CWD. This test pins the hardened ``_create_file`` check:
only the real joined target ``test_directory/test1.csv`` is a symlink, so
the write must be refused there even though the bare ``name`` is not.
"""
_create_fake_imap(mock_imaplib, with_mail=True, attachment_name="test1.csv")

# Bare ``name`` ("test1.csv") is not a symlink; the real joined target
# ("test_directory/test1.csv") is.
mock_islink.side_effect = lambda path: path == "test_directory/test1.csv"

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(name="test1.csv", local_output_directory="test_directory")

# The hardened ``_create_file`` symlink check must reject the real target
# before any file descriptor is opened, so no write goes through the link.
mock_islink.assert_any_call("test_directory/test1.csv")
mock_os_open.assert_not_called()
mock_fdopen.return_value.write.assert_not_called()

@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_uses_o_nofollow(self, mock_imaplib, mock_os_open, mock_fdopen):
"""``_create_file`` must open with ``O_NOFOLLOW`` (where available) to defeat the TOCTOU race."""
_create_fake_imap(mock_imaplib, with_mail=True, attachment_name="test1.csv")

with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(name="test1.csv", local_output_directory="test_directory")

mock_os_open.assert_called_once()
flags = mock_os_open.call_args[0][1]
# Behaviour-preserving create/truncate flags are retained (no O_EXCL).
assert flags & os.O_CREAT
assert flags & os.O_TRUNC
# O_NOFOLLOW is platform-gated; assert it is set wherever the platform has it.
if hasattr(os, "O_NOFOLLOW"):
assert flags & os.O_NOFOLLOW

@patch(open_string, new_callable=mock_open)
@patch(os_fdopen_string, new_callable=mock_open)
@patch(os_open_string)
@patch(imaplib_string)
def test_download_mail_attachments_with_mail_filter(self, mock_imaplib, mock_open_method):
def test_download_mail_attachments_with_mail_filter(self, mock_imaplib, mock_os_open, mock_fdopen):
_create_fake_imap(mock_imaplib, with_mail=True)
mail_filter = '(SINCE "01-Jan-2019")'

Expand All @@ -406,7 +471,7 @@ def test_download_mail_attachments_with_mail_filter(self, mock_imaplib, mock_ope
)

mock_imaplib.IMAP4_SSL.return_value.search.assert_called_once_with(None, mail_filter)
assert mock_open_method.call_count == 1
assert mock_os_open.call_count == 1

@patch(imaplib_string)
def test_retrieve_mail_attachments_with_max_mails(self, mock_imaplib):
Expand Down
Loading