← Back to branch summary

~dobey/ubuntuone-dev-tools/updates-2992

« back to all changes in this revision

Viewing changes to ubuntuone/devtools/services/squid.py

  • Committer: Rodney Dawes
  • Date: 2012-01-17 21:28:43 UTC
  • mfrom: (50.1.4 ubuntuone-dev-tools)
  • Revision ID: rodney.dawes@canonical.com-20120117212843-cekwgbz8k1ww57j1
Merge updates from trunk for 2.99.2 release

Show diffs side-by-side

added added

removed removed

Lines of Context:
ubuntuone/devtools/services/squid.py
 
1
#
 
2
# Copyright 2011 Canonical Ltd.
 
3
#
 
4
# This program is free software: you can redistribute it and/or modify it
 
5
# under the terms of the GNU General Public License version 3, as published
 
6
# by the Free Software Foundation.
 
7
#
 
8
# This program is distributed in the hope that it will be useful, but
 
9
# WITHOUT ANY WARRANTY; without even the implied warranties of
 
10
# MERCHANTABILITY, SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR
 
11
# PURPOSE.  See the GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License along
 
14
# with this program.  If not, see <http://www.gnu.org/licenses/>.
 
15
"""Utilities for finding and running a squid proxy for testing."""
 
16
 
 
17
import random
 
18
import signal
 
19
# pylint:disable=W0402
 
20
import string
 
21
# pylint:enable=W0402
 
22
import subprocess
 
23
import time
 
24
 
 
25
from json import dumps, loads
 
26
from os import environ, makedirs, kill, unlink
 
27
from os.path import abspath, exists, join
 
28
 
 
29
from distutils.spawn import find_executable
 
30
 
 
31
from ubuntuone.devtools.services import (
 
32
    find_config_file,
 
33
    get_arbitrary_port,
 
34
)
 
35
 
 
36
SQUID_CONFIG_FILE = 'squid.conf.in'
 
37
SQUID_DIR = 'squid'
 
38
AUTH_PROCESS_PATH = '/usr/lib/%s/ncsa_auth'
 
39
SPOOL_DIR = 'spool'
 
40
AUTH_FILE = 'htpasswd'
 
41
PROXY_ENV_VAR = 'SQUID_PROXY_SETTINGS'
 
42
 
 
43
 
 
44
def get_squid_executable():
 
45
    """Return the squid executable of the system."""
 
46
    # try with squid and if not present try with squid3 for newer systems
 
47
    # (Ubuntu P). We also return the path to the auth process so that we can
 
48
    # point to the correct one.
 
49
    squid = find_executable('squid3')
 
50
    auth_process = AUTH_PROCESS_PATH % 'squid3'
 
51
    if squid is None:
 
52
        squid = find_executable('squid')
 
53
        auth_process = AUTH_PROCESS_PATH % 'squid'
 
54
    return squid, auth_process
 
55
 
 
56
 
 
57
def get_htpasswd_executable():
 
58
    """Return the htpasswd executable."""
 
59
    return find_executable('htpasswd')
 
60
 
 
61
 
 
62
def _make_random_string(count):
 
63
    """Make a random string of the given length."""
 
64
    entropy = random.SystemRandom()
 
65
    return ''.join([entropy.choice(string.letters) for _ in
 
66
                   range(count)])
 
67
 
 
68
 
 
69
def _get_basedir(tempdir):
 
70
    """Return the base squid config."""
 
71
    basedir = join(tempdir, SQUID_DIR)
 
72
    basedir = abspath(basedir)
 
73
    if not exists(basedir):
 
74
        makedirs(basedir)
 
75
    return basedir
 
76
 
 
77
 
 
78
def _get_spool_temp_path(tempdir=''):
 
79
    """Return the temp dir to be used for spool."""
 
80
    basedir = _get_basedir(tempdir)
 
81
    path = join(basedir, SPOOL_DIR)
 
82
    path = abspath(path)
 
83
    if not exists(path):
 
84
        makedirs(path)
 
85
    return path
 
86
 
 
87
 
 
88
def _get_squid_temp_path(tempdir=''):
 
89
    """Return the temp dir to be used by squid."""
 
90
    basedir = _get_basedir(tempdir)
 
91
    path = join(basedir, SQUID_DIR)
 
92
    path = abspath(path)
 
93
    if not exists(path):
 
94
        makedirs(path)
 
95
    return path
 
96
 
 
97
 
 
98
def _get_auth_temp_path(tempdir=''):
 
99
    """Return the path for the auth file."""
 
100
    basedir = _get_basedir(tempdir)
 
101
    auth_file = join(basedir, AUTH_FILE)
 
102
    if not exists(basedir):
 
103
        makedirs(basedir)
 
104
    return auth_file
 
105
 
 
106
 
 
107
def store_proxy_settings(settings):
 
108
    """Store the proxy setting in an env var."""
 
109
    environ[PROXY_ENV_VAR] = dumps(settings)
 
110
 
 
111
 
 
112
def retrieve_proxy_settings():
 
113
    """Return the proxy settings of the env."""
 
114
    if PROXY_ENV_VAR in environ:
 
115
        return loads(environ[PROXY_ENV_VAR])
 
116
    return None
 
117
 
 
118
 
 
119
def delete_proxy_settings():
 
120
    """Delete the proxy env settings."""
 
121
    if PROXY_ENV_VAR in environ:
 
122
        del environ[PROXY_ENV_VAR]
 
123
 
 
124
 
 
125
class SquidLaunchError(Exception):
 
126
    """Error while launching squid."""
 
127
 
 
128
 
 
129
class SquidRunner(object):
 
130
    """Class for running a squid proxy with the local config."""
 
131
 
 
132
    def __init__(self):
 
133
        """Create a new instance."""
 
134
        self.squid, self.auth_process = get_squid_executable()
 
135
        if self.squid is None:
 
136
            raise SquidLaunchError('Could not locate "squid".')
 
137
 
 
138
        self.htpasswd = get_htpasswd_executable()
 
139
        if self.htpasswd is None:
 
140
            raise SquidLaunchError('Could not locate "htpasswd".')
 
141
 
 
142
        self.settings = dict(noauth_port=None, auth_port=None,
 
143
                             username=None, password=None)
 
144
        self.squid_pid = None
 
145
        self.running = False
 
146
        self.config_file = None
 
147
        self.auth_file = None
 
148
 
 
149
    def _generate_config_file(self, tempdir=''):
 
150
        """Find the first appropiate squid.conf to use."""
 
151
        # load the config file
 
152
        path = find_config_file(SQUID_CONFIG_FILE)
 
153
        # replace config settings
 
154
        basedir = join(tempdir, 'squid')
 
155
        basedir = abspath(basedir)
 
156
        if not exists(basedir):
 
157
            makedirs(basedir)
 
158
        self.config_file = join(basedir, 'squid.conf')
 
159
        with open(path) as in_file:
 
160
            template = string.Template(in_file.read())
 
161
 
 
162
        self.settings['noauth_port'] = get_arbitrary_port()
 
163
        self.settings['auth_port'] = get_arbitrary_port()
 
164
        spool_path = _get_spool_temp_path(tempdir)
 
165
        squid_path = _get_squid_temp_path(tempdir)
 
166
        with open(self.config_file, 'w') as out_file:
 
167
            out_file.write(template.safe_substitute(
 
168
                              auth_file=self.auth_file,
 
169
                              auth_process=self.auth_process,
 
170
                              noauth_port_number=self.settings['noauth_port'],
 
171
                              auth_port_number=self.settings['auth_port'],
 
172
                              spool_temp=spool_path,
 
173
                              squid_temp=squid_path))
 
174
 
 
175
    def _generate_swap(self, config_file):
 
176
        """Generate the squid swap files."""
 
177
        squid_args = ['-z', '-f', config_file]
 
178
        sp = subprocess.Popen([self.squid] + squid_args,
 
179
                              stdout=subprocess.PIPE,
 
180
                              stderr=subprocess.PIPE)
 
181
        sp.wait()
 
182
 
 
183
    def _generate_auth_file(self, tempdir=''):
 
184
        """Generates a auth file using htpasswd."""
 
185
        if self.settings['username'] is None:
 
186
            self.settings['username'] = _make_random_string(10)
 
187
        if self.settings['password'] is None:
 
188
            self.settings['password'] = _make_random_string(10)
 
189
 
 
190
        self.auth_file = _get_auth_temp_path(tempdir)
 
191
        # remove possible old auth file
 
192
        if exists(self.auth_file):
 
193
            unlink(self.auth_file)
 
194
        # create a new htpasswrd
 
195
        htpasswd_args = ['-bc',
 
196
                         self.auth_file,
 
197
                         self.settings['username'],
 
198
                         self.settings['password']]
 
199
        sp = subprocess.Popen([self.htpasswd] + htpasswd_args,
 
200
                              stdout=subprocess.PIPE,
 
201
                              stderr=subprocess.PIPE)
 
202
        sp.wait()
 
203
 
 
204
    def _is_squid_running(self):
 
205
        """Return if squid is running."""
 
206
        squid_args = ['-k', 'check', '-f', self.config_file]
 
207
        print 'Starting squid version...'
 
208
        message = 'Waiting for squid to start...'
 
209
        for timeout in (0.4, 0.1, 0.1, 0.2, 0.5, 1, 3, 5):
 
210
            try:
 
211
                #  Do not use stdout=PIPE or stderr=PIPE with this function.
 
212
                subprocess.check_call([self.squid] + squid_args,
 
213
                                       stdout=subprocess.PIPE,
 
214
                                       stderr=subprocess.PIPE)
 
215
                return True
 
216
            except subprocess.CalledProcessError:
 
217
                message += '.'
 
218
                print message
 
219
                time.sleep(timeout)
 
220
        return False
 
221
 
 
222
    def start_service(self, tempdir=None):
 
223
        """Start our own proxy."""
 
224
        # generate auth, config and swap dirs
 
225
        self._generate_auth_file(tempdir)
 
226
        self._generate_config_file(tempdir)
 
227
        self._generate_swap(self.config_file)
 
228
        squid_args = ['-N', '-X', '-f', self.config_file]
 
229
        sp = subprocess.Popen([self.squid] + squid_args,
 
230
                              stdout=subprocess.PIPE,
 
231
                              stderr=subprocess.PIPE)
 
232
        store_proxy_settings(self.settings)
 
233
        if not self._is_squid_running():
 
234
            raise SquidLaunchError('Could not start squid.')
 
235
        self.squid_pid = sp.pid
 
236
        self.running = True
 
237
 
 
238
    def stop_service(self):
 
239
        """Stop our proxy,"""
 
240
        kill(self.squid_pid, signal.SIGKILL)
 
241
        delete_proxy_settings()
 
242
        self.running = False
 
243
        unlink(self.config_file)
 
244
        unlink(self.auth_file)
 
245
        self.config_file = None