from __future__ import absolute_import

import errno
import os
import sys
import signal

from celery import _find_option_with_arg
from celery import platforms
from celery.five import open_fqdn
from celery.platforms import (
    get_fdmax,
    ignore_errno,
    set_process_title,
    signals,
    maybe_drop_privileges,
    setuid,
    setgid,
    initgroups,
    parse_uid,
    parse_gid,
    detached,
    DaemonContext,
    create_pidlock,
    Pidfile,
    LockFailed,
    setgroups,
    _setgroups_hack,
    close_open_fds,
)

try:
    import resource
except ImportError:  # pragma: no cover
    resource = None  # noqa

from celery.tests.case import (
    Case, WhateverIO, Mock, SkipTest,
    call, override_stdouts, mock_open, patch,
)


class test_find_option_with_arg(Case):

    def test_long_opt(self):
        self.assertEqual(
            _find_option_with_arg(['--foo=bar'], long_opts=['--foo']),
            'bar'
        )

    def test_short_opt(self):
        self.assertEqual(
            _find_option_with_arg(['-f', 'bar'], short_opts=['-f']),
            'bar'
        )


class test_close_open_fds(Case):

    def test_closes(self):
        with patch('os.close') as _close:
            with patch('os.closerange', create=True) as closerange:
                with patch('celery.platforms.get_fdmax') as fdmax:
                    fdmax.return_value = 3
                    close_open_fds()
                    if not closerange.called:
                        _close.assert_has_calls([call(2), call(1), call(0)])
                        _close.side_effect = OSError()
                        _close.side_effect.errno = errno.EBADF
                    close_open_fds()


class test_ignore_errno(Case):

    def test_raises_EBADF(self):
        with ignore_errno('EBADF'):
            exc = OSError()
            exc.errno = errno.EBADF
            raise exc

    def test_otherwise(self):
        with self.assertRaises(OSError):
            with ignore_errno('EBADF'):
                exc = OSError()
                exc.errno = errno.ENOENT
                raise exc


class test_set_process_title(Case):

    def when_no_setps(self):
        prev = platforms._setproctitle = platforms._setproctitle, None
        try:
            set_process_title('foo')
        finally:
            platforms._setproctitle = prev


class test_Signals(Case):

    @patch('signal.getsignal')
    def test_getitem(self, getsignal):
        signals['SIGINT']
        getsignal.assert_called_with(signal.SIGINT)

    def test_supported(self):
        self.assertTrue(signals.supported('INT'))
        self.assertFalse(signals.supported('SIGIMAGINARY'))

    def test_reset_alarm(self):
        if sys.platform == 'win32':
            raise SkipTest('signal.alarm not available on Windows')
        with patch('signal.alarm') as _alarm:
            signals.reset_alarm()
            _alarm.assert_called_with(0)

    def test_arm_alarm(self):
        if hasattr(signal, 'setitimer'):
            with patch('signal.setitimer', create=True) as seti:
                signals.arm_alarm(30)
                self.assertTrue(seti.called)

    def test_signum(self):
        self.assertEqual(signals.signum(13), 13)
        self.assertEqual(signals.signum('INT'), signal.SIGINT)
        self.assertEqual(signals.signum('SIGINT'), signal.SIGINT)
        with self.assertRaises(TypeError):
            signals.signum('int')
            signals.signum(object())

    @patch('signal.signal')
    def test_ignore(self, set):
        signals.ignore('SIGINT')
        set.assert_called_with(signals.signum('INT'), signals.ignored)
        signals.ignore('SIGTERM')
        set.assert_called_with(signals.signum('TERM'), signals.ignored)

    @patch('signal.signal')
    def test_setitem(self, set):
        handle = lambda *a: a
        signals['INT'] = handle
        set.assert_called_with(signal.SIGINT, handle)

    @patch('signal.signal')
    def test_setitem_raises(self, set):
        set.side_effect = ValueError()
        signals['INT'] = lambda *a: a


if not platforms.IS_WINDOWS:

    class test_get_fdmax(Case):

        @patch('resource.getrlimit')
        def test_when_infinity(self, getrlimit):
            with patch('os.sysconf') as sysconfig:
                sysconfig.side_effect = KeyError()
                getrlimit.return_value = [None, resource.RLIM_INFINITY]
                default = object()
                self.assertIs(get_fdmax(default), default)

        @patch('resource.getrlimit')
        def test_when_actual(self, getrlimit):
            with patch('os.sysconf') as sysconfig:
                sysconfig.side_effect = KeyError()
                getrlimit.return_value = [None, 13]
                self.assertEqual(get_fdmax(None), 13)

    class test_maybe_drop_privileges(Case):

        @patch('celery.platforms.parse_uid')
        @patch('pwd.getpwuid')
        @patch('celery.platforms.setgid')
        @patch('celery.platforms.setuid')
        @patch('celery.platforms.initgroups')
        def test_with_uid(self, initgroups, setuid, setgid,
                          getpwuid, parse_uid):

            class pw_struct(object):
                pw_gid = 50001

            def raise_on_second_call(*args, **kwargs):
                setuid.side_effect = OSError()
                setuid.side_effect.errno = errno.EPERM
            setuid.side_effect = raise_on_second_call
            getpwuid.return_value = pw_struct()
            parse_uid.return_value = 5001
            maybe_drop_privileges(uid='user')
            parse_uid.assert_called_with('user')
            getpwuid.assert_called_with(5001)
            setgid.assert_called_with(50001)
            initgroups.assert_called_with(5001, 50001)
            setuid.assert_has_calls([call(5001), call(0)])

        @patch('celery.platforms.parse_uid')
        @patch('celery.platforms.parse_gid')
        @patch('celery.platforms.setgid')
        @patch('celery.platforms.setuid')
        @patch('celery.platforms.initgroups')
        def test_with_guid(self, initgroups, setuid, setgid,
                           parse_gid, parse_uid):

            def raise_on_second_call(*args, **kwargs):
                setuid.side_effect = OSError()
                setuid.side_effect.errno = errno.EPERM
            setuid.side_effect = raise_on_second_call
            parse_uid.return_value = 5001
            parse_gid.return_value = 50001
            maybe_drop_privileges(uid='user', gid='group')
            parse_uid.assert_called_with('user')
            parse_gid.assert_called_with('group')
            setgid.assert_called_with(50001)
            initgroups.assert_called_with(5001, 50001)
            setuid.assert_has_calls([call(5001), call(0)])

            setuid.side_effect = None
            with self.assertRaises(RuntimeError):
                maybe_drop_privileges(uid='user', gid='group')
            setuid.side_effect = OSError()
            setuid.side_effect.errno = errno.EINVAL
            with self.assertRaises(OSError):
                maybe_drop_privileges(uid='user', gid='group')

        @patch('celery.platforms.setuid')
        @patch('celery.platforms.setgid')
        @patch('celery.platforms.parse_gid')
        def test_only_gid(self, parse_gid, setgid, setuid):
            parse_gid.return_value = 50001
            maybe_drop_privileges(gid='group')
            parse_gid.assert_called_with('group')
            setgid.assert_called_with(50001)
            self.assertFalse(setuid.called)

    class test_setget_uid_gid(Case):

        @patch('celery.platforms.parse_uid')
        @patch('os.setuid')
        def test_setuid(self, _setuid, parse_uid):
            parse_uid.return_value = 5001
            setuid('user')
            parse_uid.assert_called_with('user')
            _setuid.assert_called_with(5001)

        @patch('celery.platforms.parse_gid')
        @patch('os.setgid')
        def test_setgid(self, _setgid, parse_gid):
            parse_gid.return_value = 50001
            setgid('group')
            parse_gid.assert_called_with('group')
            _setgid.assert_called_with(50001)

        def test_parse_uid_when_int(self):
            self.assertEqual(parse_uid(5001), 5001)

        @patch('pwd.getpwnam')
        def test_parse_uid_when_existing_name(self, getpwnam):

            class pwent(object):
                pw_uid = 5001

            getpwnam.return_value = pwent()
            self.assertEqual(parse_uid('user'), 5001)

        @patch('pwd.getpwnam')
        def test_parse_uid_when_nonexisting_name(self, getpwnam):
            getpwnam.side_effect = KeyError('user')

            with self.assertRaises(KeyError):
                parse_uid('user')

        def test_parse_gid_when_int(self):
            self.assertEqual(parse_gid(50001), 50001)

        @patch('grp.getgrnam')
        def test_parse_gid_when_existing_name(self, getgrnam):

            class grent(object):
                gr_gid = 50001

            getgrnam.return_value = grent()
            self.assertEqual(parse_gid('group'), 50001)

        @patch('grp.getgrnam')
        def test_parse_gid_when_nonexisting_name(self, getgrnam):
            getgrnam.side_effect = KeyError('group')

            with self.assertRaises(KeyError):
                parse_gid('group')

    class test_initgroups(Case):

        @patch('pwd.getpwuid')
        @patch('os.initgroups', create=True)
        def test_with_initgroups(self, initgroups_, getpwuid):
            getpwuid.return_value = ['user']
            initgroups(5001, 50001)
            initgroups_.assert_called_with('user', 50001)

        @patch('celery.platforms.setgroups')
        @patch('grp.getgrall')
        @patch('pwd.getpwuid')
        def test_without_initgroups(self, getpwuid, getgrall, setgroups):
            prev = getattr(os, 'initgroups', None)
            try:
                delattr(os, 'initgroups')
            except AttributeError:
                pass
            try:
                getpwuid.return_value = ['user']

                class grent(object):
                    gr_mem = ['user']

                    def __init__(self, gid):
                        self.gr_gid = gid

                getgrall.return_value = [grent(1), grent(2), grent(3)]
                initgroups(5001, 50001)
                setgroups.assert_called_with([1, 2, 3])
            finally:
                if prev:
                    os.initgroups = prev

    class test_detached(Case):

        def test_without_resource(self):
            prev, platforms.resource = platforms.resource, None
            try:
                with self.assertRaises(RuntimeError):
                    detached()
            finally:
                platforms.resource = prev

        @patch('celery.platforms._create_pidlock')
        @patch('celery.platforms.signals')
        @patch('celery.platforms.maybe_drop_privileges')
        @patch('os.geteuid')
        @patch(open_fqdn)
        def test_default(self, open, geteuid, maybe_drop,
                         signals, pidlock):
            geteuid.return_value = 0
            context = detached(uid='user', gid='group')
            self.assertIsInstance(context, DaemonContext)
            signals.reset.assert_called_with('SIGCLD')
            maybe_drop.assert_called_with(uid='user', gid='group')
            open.return_value = Mock()

            geteuid.return_value = 5001
            context = detached(uid='user', gid='group', logfile='/foo/bar')
            self.assertIsInstance(context, DaemonContext)
            self.assertTrue(context.after_chdir)
            context.after_chdir()
            open.assert_called_with('/foo/bar', 'a')
            open.return_value.close.assert_called_with()

            context = detached(pidfile='/foo/bar/pid')
            self.assertIsInstance(context, DaemonContext)
            self.assertTrue(context.after_chdir)
            context.after_chdir()
            pidlock.assert_called_with('/foo/bar/pid')

    class test_DaemonContext(Case):

        @patch('os.fork')
        @patch('os.setsid')
        @patch('os._exit')
        @patch('os.chdir')
        @patch('os.umask')
        @patch('os.close')
        @patch('os.closerange')
        @patch('os.open')
        @patch('os.dup2')
        def test_open(self, dup2, open, close, closer, umask, chdir,
                      _exit, setsid, fork):
            x = DaemonContext(workdir='/opt/workdir', umask=0o22)
            x.stdfds = [0, 1, 2]

            fork.return_value = 0
            with x:
                self.assertTrue(x._is_open)
                with x:
                    pass
            self.assertEqual(fork.call_count, 2)
            setsid.assert_called_with()
            self.assertFalse(_exit.called)

            chdir.assert_called_with(x.workdir)
            umask.assert_called_with(0o22)
            self.assertTrue(dup2.called)

            fork.reset_mock()
            fork.return_value = 1
            x = DaemonContext(workdir='/opt/workdir')
            x.stdfds = [0, 1, 2]
            with x:
                pass
            self.assertEqual(fork.call_count, 1)
            _exit.assert_called_with(0)

            x = DaemonContext(workdir='/opt/workdir', fake=True)
            x.stdfds = [0, 1, 2]
            x._detach = Mock()
            with x:
                pass
            self.assertFalse(x._detach.called)

            x.after_chdir = Mock()
            with x:
                pass
            x.after_chdir.assert_called_with()

    class test_Pidfile(Case):

        @patch('celery.platforms.Pidfile')
        def test_create_pidlock(self, Pidfile):
            p = Pidfile.return_value = Mock()
            p.is_locked.return_value = True
            p.remove_if_stale.return_value = False
            with override_stdouts() as (_, err):
                with self.assertRaises(SystemExit):