1# A minimal client for Mercurial's command server
2
3from __future__ import absolute_import, print_function
4
5import io
6import os
7import re
8import signal
9import socket
10import struct
11import subprocess
12import sys
13import time
14
15if sys.version_info[0] >= 3:
16    stdout = sys.stdout.buffer
17    stderr = sys.stderr.buffer
18    stringio = io.BytesIO
19
20    def bprint(*args):
21        # remove b'' as well for ease of test migration
22        pargs = [re.sub(br'''\bb(['"])''', br'\1', b'%s' % a) for a in args]
23        stdout.write(b' '.join(pargs) + b'\n')
24
25
26else:
27    import cStringIO
28
29    stdout = sys.stdout
30    stderr = sys.stderr
31    stringio = cStringIO.StringIO
32    bprint = print
33
34
35def connectpipe(path=None, extraargs=()):
36    cmdline = [b'hg', b'serve', b'--cmdserver', b'pipe']
37    if path:
38        cmdline += [b'-R', path]
39    cmdline.extend(extraargs)
40
41    def tonative(cmdline):
42        if os.name != 'nt':
43            return cmdline
44        return [arg.decode("utf-8") for arg in cmdline]
45
46    server = subprocess.Popen(
47        tonative(cmdline), stdin=subprocess.PIPE, stdout=subprocess.PIPE
48    )
49
50    return server
51
52
53class unixconnection(object):
54    def __init__(self, sockpath):
55        self.sock = sock = socket.socket(socket.AF_UNIX)
56        sock.connect(sockpath)
57        self.stdin = sock.makefile('wb')
58        self.stdout = sock.makefile('rb')
59
60    def wait(self):
61        self.stdin.close()
62        self.stdout.close()
63        self.sock.close()
64
65
66class unixserver(object):
67    def __init__(self, sockpath, logpath=None, repopath=None):
68        self.sockpath = sockpath
69        cmdline = [b'hg', b'serve', b'--cmdserver', b'unix', b'-a', sockpath]
70        if repopath:
71            cmdline += [b'-R', repopath]
72        if logpath:
73            stdout = open(logpath, 'a')
74            stderr = subprocess.STDOUT
75        else:
76            stdout = stderr = None
77        self.server = subprocess.Popen(cmdline, stdout=stdout, stderr=stderr)
78        # wait for listen()
79        while self.server.poll() is None:
80            if os.path.exists(sockpath):
81                break
82            time.sleep(0.1)
83
84    def connect(self):
85        return unixconnection(self.sockpath)
86
87    def shutdown(self):
88        os.kill(self.server.pid, signal.SIGTERM)
89        self.server.wait()
90
91
92def writeblock(server, data):
93    server.stdin.write(struct.pack(b'>I', len(data)))
94    server.stdin.write(data)
95    server.stdin.flush()
96
97
98def readchannel(server):
99    data = server.stdout.read(5)
100    if not data:
101        raise EOFError
102    channel, length = struct.unpack('>cI', data)
103    if channel in b'IL':
104        return channel, length
105    else:
106        return channel, server.stdout.read(length)
107
108
109def sep(text):
110    return text.replace(b'\\', b'/')
111
112
113def runcommand(
114    server, args, output=stdout, error=stderr, input=None, outfilter=lambda x: x
115):
116    bprint(b'*** runcommand', b' '.join(args))
117    stdout.flush()
118    server.stdin.write(b'runcommand\n')
119    writeblock(server, b'\0'.join(args))
120
121    if not input:
122        input = stringio()
123
124    while True:
125        ch, data = readchannel(server)
126        if ch == b'o':
127            output.write(outfilter(data))
128            output.flush()
129        elif ch == b'e':
130            error.write(data)
131            error.flush()
132        elif ch == b'I':
133            writeblock(server, input.read(data))
134        elif ch == b'L':
135            writeblock(server, input.readline(data))
136        elif ch == b'm':
137            bprint(b"message: %r" % data)
138        elif ch == b'r':
139            (ret,) = struct.unpack('>i', data)
140            if ret != 0:
141                bprint(b' [%d]' % ret)
142            return ret
143        else:
144            bprint(b"unexpected channel %c: %r" % (ch, data))
145            if ch.isupper():
146                return
147
148
149def check(func, connect=connectpipe):
150    stdout.flush()
151    server = connect()
152    try:
153        return func(server)
154    finally:
155        server.stdin.close()
156        server.wait()
157
158
159def checkwith(connect=connectpipe, **kwargs):
160    def wrap(func):
161        return check(func, lambda: connect(**kwargs))
162
163    return wrap
164