1

"""

     2

    A simple TCP client in the kind of syslog.fifo/file.

     3

     4

    Interface: fileno(), __iter__, __call__

     5

"""

     6

     7

# XXX: absolute import plz

     8

socket = __import__('socket')

     9

    10

import select

    11

import errno

    12

    13

import urlparse

    14

    15

import logging; log = logging.getLogger('pvl.socket')

    16

    17

# order matters!

    18

URL = (

    19

    # scheme    family              socktype

    20

    ( 'unix',   (socket.AF_UNIX,    None                )   ), # socktype is given

    21

    ( 'tcp',    (0,                 socket.SOCK_STREAM  )   ), # AF_UNSPEC

    22

    ( 'udp',    (0,                 socket.SOCK_DGRAM   )   ), # AF_UNSPEC

    23

)

    24

    25

URL_SCHEMES = dict(URL)

    26

    27

def parse (str, port=None, scheme='tcp', unix=socket.SOCK_DGRAM) :

    28

    """

    29

        Parse given string into (AF_*, SOCK_*, host, port).

    30

    31

        For AF_UNIX, the path is in host, and port is empty, and the socktype is the given unix=... value.

    32

    """

    33

    34

    family, socktype = URL_SCHEMES[scheme]

    35

    url = urlparse.urlparse(str)

    36

    37

    # TODO: UNIX?

    38

    if url.scheme and url.netloc :

    39

        # proper url

    40

        family, socktype = URL_SCHEMES[url.scheme]

    41

    42

        return family, socktype, url.hostname, url.port or port

    43

    44

    elif url.scheme and url.path :

    45

        # host:port

    46

        return family, socktype, url.scheme, int(url.path)

    47

    48

    elif url.path :

    49

        # host

    50

        return family, socktype, url.path, port

    51

    52

    else :

    53

        raise ValueError("unparseable connect URL: %s", str)

    54

    55

def connect (str, *args, **kwargs) :

    56

    """

    57

        Returns a connected socket for given parse()'d string.

    58

    """

    59

    60

    family, socktype, host, port = parse(str, *args, **kwargs)

    61

    62

    if family == socket.AF_UNIX :

    63

        raise ValueError("XXX: AF_UNIX is not yet supported", str)

    64

    65

    else : # AF_UNSPEC

    66

        return connect_inet(host, port, family=family, socktype=socktype)

    67

    68

def connect_inet (host=None, port=None, family=socket.AF_UNSPEC, socktype=socket.SOCK_STREAM) :

    69

    """

    70

        Return a TCP/UDP socket connected to the given host/port using getaddrinfo.

    71

    72

        TODO: timeout?

    73

    """

    74

    75

    log.debug("%s:%s: family=%s, socktype=%s", host, port, family, socktype)

    76

    77

    if host :

    78

        flags = socket.AI_CANONNAME

    79

    else :

    80

        flags = 0

    81

    82

    addrinfo = socket.getaddrinfo(host, port, family, socktype, 0, flags)

    83

    84

    if not addrinfo :

    85

        raise Exception("getaddrinfo: %s:%s: no results" % (host, port))

    86

    87

    for af, st, proto, name, addr in addrinfo :

    88

        try :

    89

            sock = socket.socket(af, st, proto)

    90

    91

        except socket.error as error :

    92

            log.warning("%s:%s: socket: %s", host, port, error)

    93

            continue

    94

    95

        log.debug("%s:%s: socket: %s", host, port, sock)

    96

    97

        try :

    98

            sock.connect(addr)

    99

   100

        except socket.error as error :

   101

            log.warning("%s:%s: connect: %s", host, port, error)

   102

            continue

   103

   104

        log.debug("%s:%s: connect", host, port)

   105

        log.info("%s", name)

   106

   107

        return sock

   108

   109

    else :

   110

        raise Exception("Unable to connect: %s:%s: %s" % (host, port, error))

   111

   112

def reverse (sockaddr, numeric_host=False, numeric_port=True) :

   113

    """

   114

        Resolve given sockaddr, returning (host, port).

   115

    """

   116

   117

    flags = 0

   118

   119

    if numeric_host :

   120

        flags |= socket.NI_NUMERICHOST

   121

   122

    if numeric_port :

   123

        flags |= socket.NI_NUMERICSERV

   124

   125

    return socket.getnameinfo(sockaddr, flags)

   126

   127

def socket_str (sock) :

   128

    peer = sock.getpeername()

   129

   130

    # lookup scheme

   131

    for scheme, (family, socktype) in URL :

   132

        if family and family != sock.family :

   133

            continue

   134

        elif socktype and socktype != sock.type :

   135

            continue

   136

        else :

   137

            break

   138

    else :

   139

        scheme = None

   140

   141

    host, port = reverse(peer)

   142

   143

    if scheme :

   144

        return "{scheme}://{host}:{port}".format(scheme=scheme, host=host, port=port)

   145

    else :

   146

        return "{host}:{port}".format(host=host, port=port)

   147

   148

# XXX: SIGPIPE

   149

def nonblocking (call, *args, **kwargs) :

   150

    """

   151

        Call the given function, which read/writes on a nonblocking file, and return None if it would have blocked.

   152

    """

   153

   154

    try :

   155

        return call(*args, **kwargs)

   156

   157

    except socket.error as ex :

   158

        # block?

   159

        if ex.errno == errno.EAGAIN or ex.errno == errno.EWOULDBLOCK:

   160

            # empty

   161

            return None

   162

   163

        else :

   164

            raise

   165

   166

class ReadStream (object) :

   167

    """

   168

        Buffered stream, supporting non-blocking/line-based reads.

   169

    """

   170

   171

    BLOCK=512

   172

   173

    def __init__ (self, sock, buffer=None) :

   174

        """

   175

            TODO: buffer    - maximum line length

   176

        """

   177

   178

        self.sock = sock

   179

        self._buf = ''

   180

   181

    def fileno (self) :

   182

        return self.sock.fileno()

   183

   184

    def _read (self, block=BLOCK) :

   185

        """

   186

            Read up to n bytes from socket.

   187

   188

            Returns None if we would block.

   189

            Raises EOFError on EOF.

   190

        """

   191

   192

        buf = nonblocking(self.sock.recv, block)

   193

   194

        log.debug("%s: %s", self, buf)

   195

   196

        if buf is None :

   197

            return None

   198

        elif buf :

   199

            return buf

   200

        else :

   201

            raise EOFError()

   202

   203

    def peek (self) :

   204

        """

   205

            Peek at data in buffer.

   206

        """

   207

   208

        return self._buf

   209

   210

    def read (self) :

   211

        """

   212

            Read and return any available input.

   213

   214

            Returns None if blocking.

   215

        """

   216

   217

        if self._buf :

   218

            buf, self._buf = self._buf, ''

   219

   220

        else :

   221

            buf = self._read()

   222

   223

        return buf

   224

   225

    def readline (self) :

   226

        """

   227

            Read and return next waiting line from input.

   228

   229

            Line is returned without trailing '\r\n' or '\n'.

   230

   231

            Returns None if there is no line available.

   234

        """

   235

   236

        while '\n' not in self._buf :

   237

            # read chunk

   238

            read = self._read()

   239

   240

            if read is None :

   241

                return None

   242

   243

            self._buf += read

   244

   245

        # split out one line

   246

        line, self._buf = self._buf.split('\n', 1)

   247

   248

        # in case we had \r\n

   249

        line = line.rstrip('\r')

   250

   251

        log.debug("%s: %s", self, line)

   252

   253

        return line

   254

   255

    def readlines (self) :

   256

        """

   257

            Read any available input, yielding lines.

   258

   259

            Returns None if thre is no more input available.

   260

   261

            Raises EOFError in the socket was closed.

   262

        """

   263

   264

        while True :

   265

            line = self.readline()

   266

   267

            if line is None :

   268

                return

   269

            else :

   270

                yield line

   271

   272

    __iter__ = readlines

   273

   274

    def __str__ (self) :

   275

        return socket_str(self.sock)

   276

   277

class WriteStream (object) :

   278

    """

   279

        Writable stream, supporting non-blocking/buffered writes.

   280

   281

        XXX: buffering is completely untested

   282

    """

   283

   284

    EOL = '\n'

   285

   286

    def __init__ (self, sock, buffer=None) :

   287

        """

   288

            TODO:   buffer  - maximum outgoing buffer length

   289

        """

   290

   291

        self.sock = sock

   292

        self._buf = buffer

   293

   294

    def _write (self, buf) :

   295

        """

   296

            Write given data to socket, returning the number of bytes written, or None, if buffering is enabled.

   297

        """

   298

   299

        send = nonblocking(self.sock.send, buf)

   300

   301

        # eof on write?

   302

        if send is None :

   303

            return None

   304

   305

        elif send :

   306

            # ok, message (partially) written

   307

            return send

   308

   309

        else :

   310

            # XXX: zero-length send? how do we handle this? What does it actually mean?

   311

            # handle as a wouldblock...

   312

            return None

   313

   314

    def write (self, data) :

   315

        """

   316

            Write given data to socket.

   317

   318

            TODO: buffer small chunks -> select writable -> write?

   319

   320

            Buffers if not able to write, or raises EOFError (hah!)

   321

        """

   322

   323

        if not self._buf :

   324

            # write directly

   325

            while data :

   326

                write = self._write(data)

   327

   328

                if write :

   329

                    # remaining data

   330

                    data = data[write:]

   331

   332

                else :

   333

                    # cannot write more

   334

                    break

   335

   336

        if not data :

   337

            # sent

   338

            return

   339

   340

        if self._buf is None :

   341

            # no write buffering, and socket buffer full!

   342

            raise EOFError()

   343

   344

        # append to outgoing buffer

   345

        self._buf += data

   346

   347

    def writeline (self, line, eol=EOL) :

   348

        """

   349

            Write out line.

   350

        """

   351

   352

        log.debug("%s: %s", self, line)

   353

   354

        self.write(str(line))

   355

        self.write(eol)

   356

   357

    def __call__ (self, *lines) :

   358

        for line in lines :

   359

            self.writeline(line)

   360

   361

        # TODO: flush

   362

   363

    def __str__ (self) :

   364

        return socket_str(self.sock)

   365

   366