1

"""

     2

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

     3

     4

    Interface: fileno(), __iter__, __call__

     5

"""

     6

     7

from __future__ import absolute_import

     8

     9

import errno

    10

import select

    11

import socket

    12

import urlparse

    13

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

    14

    15

from socket import SOCK_STREAM

    16

    17

def parse (url, port=None) :

    18

    """

    19

        Parse given string into host, port, path parts.

    20

    21

        >>> parse_url('')

    22

        (None, None, None)

    23

        >>> parse_url('foo')

    24

        ('foo', None, None)

    25

        >>> parse_url('foo:80')

    26

        ('foo', 80, None)

    27

        >>> parse_url('foo:http')

    28

        ('foo', 'http', None)

    29

        >>> parse_url('/run/foo.sock')

    30

        (None, None, 'run/foo.sock')

    31

    """

    32

    33

    if '/' in url :

    34

        url, path = url.split('/', 1)

    35

    else :

    36

        path = None

    37

    38

    if ':' in url :

    39

        url, port = url.split(':')

    40

    41

        if port.isdigit() :

    42

            port = int(port)

    43

    44

    if url :

    45

        host = url

    46

    else :

    47

        host = None

    48

    49

    return host, port, path

    50

    51

def connect (socktype, url, port=None, family=None) :

    52

    """

    53

        Returns a connected socket for given string, based on parse().

    54

    """

    55

    56

    host, port, path = parse(url, port=port)

    57

    58

    # autodetect as AF_UNIX

    59

    if path and not family :

    60

        family = socket.AF_UNIX

    61

    62

    if family == socket.AF_UNIX :

    63

        raise ValueError("TODO: AF_UNIX is not yet supported: %s" % (url, ))

    64

    65

    else : # AF_UNSPEC or AF_INET/AF_INET6

    66

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

    67

    68

def connect_inet (socktype, host, port, family=None) :

    69

    """

    70

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

    71

    72

        socktype: SOCK_STREAM or SOCK_DGRAM

    73

        host: IP address, hostname, or None for localhost.

    74

        port: integer port, or string service.

    75

        family: AF_UNSPEC for IP/DNS dependent IPv6/4, or AF_INET or AF_INET6.

    76

    77

        TODO: timeout?

    78

    """

    79

    80

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

    81

    82

    if family is None :

    83

        family = socket.AF_UNSPEC

    84

    85

    if host :

    86

        flags = socket.AI_CANONNAME

    87

    else :

    88

        flags = 0

    89

    90

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

    91

    92

    if not addrinfo :

    93

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

    94

    95

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

    96

        try :

    97

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

    98

    99

        except socket.error as error :

   100

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

   101

            continue

   102

   103

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

   104

   105

        try :

   106

            sock.connect(addr)

   107

   108

        except socket.error as error :

   109

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

   110

            continue

   111

   112

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

   113

   114

        return sock

   115

   116

    else :

   117

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

   118

   119

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

   120

    """

   121

        Resolve given sockaddr, returning (host, port).

   122

    """

   123

   124

    flags = 0

   125

   126

    if numeric_host :

   127

        flags |= socket.NI_NUMERICHOST

   128

   129

    if numeric_port :

   130

        flags |= socket.NI_NUMERICSERV

   131

   132

    return socket.getnameinfo(sockaddr, flags)

   133

   134

def socket_str (sock) :

   135

    """

   136

        Return url str for socket peer.

   137

    """

   138

   139

    # get connected peer

   140

    try :

   141

        peer = sock.getpeername()

   142

   143

    except socket.error as ex :

   144

        # fails if socket is not connected XXX: even after EOF on read..?

   145

        return str(ex)

   146

   147

    host, port = reverse(peer)

   148

   149

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

   150

   151

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

   152

    """

   153

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

   154

   155

        Raises EOFError on SIGPIPE/EPIPE.

   156

   157

        # XXX: does python handle SIGPIPE for us?

   158

    """

   159

   160

    try :

   161

        return call(*args, **kwargs)

   162

   163

    except socket.error as ex :

   164

        # block?

   165

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

   166

            # empty

   167

            return None

   168

   169

        elif ex.errno == errno.EPIPE :

   170

            # XXX: write-eof?

   171

            raise EOFError()

   172

   173

        else :

   174

            raise

   175

   176

class ReadStream (object) :

   177

    """

   178

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

   179

    """

   180

   181

    BLOCK = 512

   182

   183

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

   184

        """

   185

            TODO: buffer    - maximum line length

   186

        """

   187

   188

        self.sock = sock

   189

        self._buf = ''

   190

   191

    def fileno (self) :

   192

        return self.sock.fileno()

   193

   194

    def _read (self, block=BLOCK) :

   195

        """

   196

            Read up to n bytes from socket.

   197

   198

            Returns None if we would block.

   199

            Raises EOFError on EOF.

   200

        """

   201

   202

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

   203

   204

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

   205

   206

        if buf is None :

   207

            return None

   208

        elif buf :

   209

            return buf

   210

        else :

   211

            raise EOFError()

   212

   213

    def peek (self) :

   214

        """

   215

            Peek at data in buffer.

   216

        """

   217

   218

        return self._buf

   219

   220

    def read (self) :

   221

        """

   222

            Read and return any available input.

   223

   224

            Returns None if blocking.

   225

        """

   226

   227

        if self._buf :

   228

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

   229

   230

        else :

   231

            buf = self._read()

   232

   233

        return buf

   234

   235

    __call__ = read

   236

   237

    def readline (self) :

   238

        """

   239

            Read and return next waiting line from input.

   240

   241

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

   242

   243

            Returns None if there is no line available.

   244

   245

            XXX: trailing data in buf when _read() raises EOFError?

   246

        """

   247

   248

        while '\n' not in self._buf :

   249

            # read chunk

   250

            read = self._read()

   251

   252

            if read is None :

   253

                return None

   254

   255

            self._buf += read

   256

   257

        # split out one line

   258

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

   259

   260

        # in case we had \r\n

   261

        line = line.rstrip('\r')

   262

   263

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

   264

   265

        return line

   266

   267

    def readlines (self) :

   268

        """

   269

            Read any available input, yielding lines.

   270

   271

            Returns None if thre is no more input available.

   272

   273

            Raises EOFError in the socket was closed.

   274

        """

   275

   276

        while True :

   277

            line = self.readline()

   278

   279

            if line is None :

   280

                return

   281

            else :

   282

                yield line

   283

   284

    __iter__ = readlines

   285

   286

    def __str__ (self) :

   287

        return socket_str(self.sock)

   288

   289

class WriteStream (object) :

   290

    """

   291

        Writable stream, supporting non-blocking/buffered writes.

   292

   293

        XXX: buffering is completely untested

   294

    """

   295

   296

    EOL = '\n'

   297

   298

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

   299

        """

   300

            TODO:   buffer  - maximum outgoing buffer length

   301

        """

   302

   303

        self.sock = sock

   304

        self._buf = buffer

   305

        self.eol = eol or self.EOL

   306

   307

    def _write (self, buf) :

   308

        """

   309

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

   310

        """

   311

   312

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

   313

   314

        # eof on write?

   315

        if send is None :

   316

            return None

   317

   318

        elif send :

   319

            # ok, message (partially) written

   320

            return send

   321

   322

        else :

   323

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

   324

            # handle as a wouldblock...

   325

            return None

   326

   327

    def write (self, data) :

   328

        """

   329

            Write given data to socket.

   330

   331

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

   332

   333

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

   334

        """

   335

   336

        if not self._buf :

   337

            # write directly

   338

            while data :

   339

                write = self._write(data)

   340

   341

                if write :

   342

                    # remaining data

   343

                    data = data[write:]

   344

   345

                else :

   346

                    # cannot write more

   347

                    break

   348

   349

        if not data :

   350

            # sent

   351

            return

   352

   353

        if self._buf is None :

   354

            # no write buffering, and socket buffer full!

   355

            raise EOFError()

   356

   357

        # append to outgoing buffer

   358

        self._buf += data

   359

   360

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

   361

        """

   362

            Write out line.

   363

        """

   364

   365

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

   366

   367

        self.write(str(line))

   368

        self.write(eol or self.eol)

   369

   370

    def __call__ (self, *lines) :

   371

        for line in lines :

   372

            self.writeline(line)

   373

   374

        # TODO: flush

   375

   376

    def __str__ (self) :

   377

        return socket_str(self.sock)

   378

   379

class SockStream (object) :

   380

    """

   381

        A (TCP) socket connection.

   382

   383

        Supports nonblocking reads and line-oriented protoocls.

   384

    """

   385

   386

    PORT = None

   387

    EOL = None

   388

   389

    @classmethod

   390

    def connect (cls, host, port=None, family=None, **opts) :

   391

        """

   392

            Blocking TCP client resolve/connect.

   393

   394

            Raises ??? if connect fails.

   395

        """

   396

   397

        sock = connect(socket.SOCK_STREAM, host, port or cls.PORT,

   398

                family=family,

   399

        )

   400

   401

        tcp = cls(sock, **opts)

   402

        log.info("%s", tcp)

   403

        return tcp

   404

   405

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

   406

        # store

   407

        self.sock = sock

   408

   409

        # nonblocking mode?

   410

        if nonblocking :

   411

            self.sock.setblocking(not nonblocking)

   412

   413

        # read/write buffer

   414

        self.read = ReadStream(sock)

   415

        self.write = WriteStream(sock, eol=self.EOL)

   416

   417

        # cached endpoint names

   418

        self._local_name = None

   419

        self._remote_name = None

   420

   421

    def fileno (self) :

   422

        return self.sock.fileno()

   423

   424

    def __iter__ (self) :

   425

        return iter(self.read)

   426

   427

    def __call__ (self, *args, **kwargs) :

   428

        return self.write(*args, **kwargs)

   429

   430

    @property

   431

    def local (self) :

   432

        """

   433

            Resolve the local endpoint (host, port)

   434

        """

   435

   436

        if not self._local_name :

   437

            try :

   438

                self._local_name, port = reverse(self.sock.getsockname())

   439

            except socket.error as ex:

   440

                return None

   441

   442

        return self._local_name

   443

   444

    @property

   445

    def remote (self) :

   446

        """

   447

            Resolve the remote endpoint (host, port).

   448

        """

   449

   450

        if not self._remote_name :

   451

            try :

   452

                self._remote_name, port = reverse(self.sock.getpeername())

   453

            except socket.error as ex:

   454

                return None

   455

   456

        return self._remote_name

   457

   458

    def __str__ (self) :

   459

        return "%s" % self.remote

   460