"""
* ufw.io.streams.sock
* ufw stream plugin
provides:
    proto sock (other=None) uses (path if object is None)
    proto unix (server=False) uses path
    proto inet (server=False) uses host, port, path
requires:
    streambase readwrite
"""

import os
import ssl
import socket
from select import select
from ufw.errors import *
from ufw.core.object import *
from ufw.io.streams.base import *

__all__ = ["sock", "inet", "unix"]

class sock (readwrite):
    flags = {
        "ioctl":True,
        "secure":True,
        "server":True
    }
    def open(self, other=None):
        self.server = False
        if not other: other = object.get(self.parent.path.path)
        self.stream = other
        self.stream.setblocking(0)
        return self.parent
    def read(self, ln=-1, bufsize=4096, wait=None, rs=''):
        self.opened()
        select([self.stream], [], [])
        if ln >= 0:
            self.opened()
            buf = bytes()
            while len(buf) < ln:
                try: buf += self.stream.recv(ln)
                except socket.error: break
        else:
            buf = bytes()
            while True:
                self.opened()
                try: add = self.stream.recv(bufsize)
                except socket.error: add = bytes()
                buf += add
                if not wait and not add: break
                if wait and wait in add: break
                elif wait:
                    select([self], [], [])
        buf = str(buf)
        self._notnull(buf)
        if wait and wait in buf:
            buf = buf[:buf.find(wait)]
            buf = buf.rstrip(rs)
            return buf + wait
        return buf
    def readline(self):
        return self.read(-1, 1, "\n", "\r\n")
    def write(self, buf):
        self.opened()
        if isinstance(buf, unicode):
            buf = buf.encode("utf-8")
        while buf:
            r, w, x = select([], [self.stream], [], 30)
            if self.stream not in w:
                self.close()
                raise io_closed
            sent = self.stream.send(buf)
            buf = buf[sent:]
    def encrypt(self):
        self.opened()
        self.ssl = True
        s = ssl.wrap_socket(
            self.stream,
            server_side=self.server,
            do_handshake_on_connect=False
        )
        ssl_may_fail(s, s.do_handshake)
        self.opened()
        self.stream = s
    def decrypt(self):
        if not self.ssl: return
        self.opened()
        self.ssl = False
        s = ssl_may_fail(self.stream, self.stream.unwrap)
        self.opened()
        self.stream = s
    def listen(self, pint=-1):
        assert self.server
        if pint < 0:
            select([self.stream], [], [])
        else:
            r, w, x = select([self.stream], [], [], pint)
            if self.stream not in r: return None, None
        return self.stream.accept()
    def fileno(self):
        return self.stream.fileno()

class unix (sock):
    def __init__(self, *a, **k):
        if not hasattr(socket, "AF_UNIX"):
            raise io_invalid_protocol ("Windows cannot into AF_UNIX")
        super(unix, self).__init__(*a, **k)
    def open(self, server=False):
        self.server = server
        self.stream = socket.socket(socket.AF_UNIX)
        if server:
            self.stream.bind(self.parent.path.path)
            self.stream.listen(5)
        else:
            self.stream.connect(self.parent.path.path)
        self.stream.setblocking(0)
        return self.parent
    def close(self):
        self.opened()
        del self.stream
        if self.server:
            os.unlink(self.parent.path.path)

class inet (sock):
    def open(self, server=False):
        self.ssl = False
        self.server = server
        addr = self.parent.path.hostname
        port = self.parent.path.port
        if not addr: addr = ""
        if not port: port = 0
        self.stream = socket.socket(socket.AF_INET)
        if server:
            if addr == "localhost": addr = ""
            self.stream.bind((addr, port))
            self.stream.listen(5)
        else:
            self.stream.connect((addr, port))
        self.stream.setblocking(0)
        return self.parent

def ssl_may_fail(s, f, *a, **k):
    while True:
        try: return f(*a, **k)
        except ssl.SSLError as e:
            if e.args[0] == ssl.SSL_ERROR_WANT_READ: select([s],[],[])
            elif e.args[0] == ssl.SSL_ERROR_WANT_WRITE: select([],[s],[])
            else: raise