~launchpad-pqm/launchpad/devel : revision 13447.1.2

1

#!/usr/bin/python

2

#

3

4

# GNU Affero General Public License version 3 (see the file LICENSE).

5

6

"""

7

Generate an EC2 image that is capable of running the Windmill browser UI

8

testing tool.

9

10

You must provide a base image that will be augmented with the necessary

11

packages and configuration.

12

13

The script requires certain options to be specified in order to function

14

properly. These options may be supplied using command-line switches, or

15

via a config file, with the --config command-line switch. The default

16

config file location is ~/.ec2/ec2bundle.cfg

17

18

The config file format simply replicates the required command-line options

19

as configuration keys.

20

21

---- ec2bundle.cfg ---

22

23

[DEFAULT]

24

key = gsg-keypair

25

identity-file = ~/.ec2/foo-keypair-id_rsa

26

private-key = ~/.ec2/pk-HKZYKTAIG2ECMXYIBH3HXV4ZBZQ55CLO.pem

27

cert =~/.ec2/cert-HKZYKTAIG2ECMXYIBH3HXV4ZBZQ55CLO.pem

28

user-id = AIDADH4IGTRXXKCD

29

access-key = SOMEBIGSTRINGOFDIGITS

30

secret-key = s0m3funKyStr1Ng0fD1gitZ

31

#bucket = foo # Required, but you probably want to customize it each time.

32

33

---- fin ---

34

35

"""

36

37

__metatype__ = type

38

39

40

# Reuse a whole bunch of code from ec2test.py.

41

import ConfigParser

42

import ec2test

43

import logging

44

import optparse

45

import os

46

import paramiko

47

import select

48

import socket

49

import subprocess

50

import sys

51

import time

52

53

54

log = logging.getLogger(__name__)

55

info = log.info

56

debug = log.debug

57

58

59

usage = """

60

Generate an EC2 image for Windmill testing in Firefox.

61

62

usage: %prog [options] AMI-ID

63

"""

64

65

class Instance:

66

"""An EC2 instance controller."""

67

68

def __init__(self, instance):

69

self._instance = instance

70

71

@property

72

def id(self):

73

return self._instance.id

74

75

@property

76

def hostname(self):

77

return self._instance.public_dns_name

78

79

def stop(self):

80

instance = self._instance

81

82

instance.update()

83

if instance.state not in ('shutting-down', 'terminated'):

84

# terminate instance

85

instance.stop()

86

instance.update()

87

88

info('instance %s\n' % (instance.state,))

89

90

def wait_for_instance_to_start(self):

91

"""Wait for the instance to transition to the "running" state."""

92

instance = self._instance

93

info('Instance %s starting..' % instance.id)

94

95

start = time.time()

96

while instance.state == 'pending':

97

sys.stdout.write('.')

98

sys.stdout.flush()

99

time.sleep(5)

100

instance.update()

101

if instance.state == 'running':

102

info('\ninstance now running at %s\n' % instance.public_dns_name)

103

elapsed = time.time() - start

104

info('Started in %d minutes %d seconds\n' %

105

(elapsed // 60, elapsed % 60))

106

cout = instance.get_console_output()

107

info(cout.output)

108

else:

109

raise RuntimeError('failed to start: %s\n' % instance.state)

110

111

@classmethod

112

def from_image(cls, account, ami_id, instance_type):

113

"""Return a new instance using the given startup parameters."""

114

info("Starting instance")

115

116

# Set up a security group that opens up ports 22, 80, and 443. Also

117

# opens up access for our IP.

118

account.acquire_security_group()

119

120

image = account.acquire_image(ami_id)

121

key = account.name

122

debug("Image: %s, Type: %s, Key: %s" % (

123

ami_id, instance_type, key))

124

125

reservation = image.run(

126

key_name=key,

127

security_groups=[key],

128

instance_type=instance_type)

129

130

instance = cls(reservation.instances[0])

131

instance.wait_for_instance_to_start()

132

return instance

133

134

@classmethod

135

def from_running_instance(cls, account, instance_id):

136

"""Create an object from an already running EC2 instance."""

137

instance = account.get_instance(instance_id)

138

if not instance:

139

raise RuntimeError(

140

"Unable to connect to instance %s" % instance_id)

141

142

info("Connected to instance %s" % instance_id)

143

proxy = cls(instance)

144

# Just to be extra safe.

145

proxy.wait_for_instance_to_start()

146

return proxy

147

148

149

class SSHConnector:

150

"""Handle the various aspects of using an SSH connection."""

151

152

def __init__(self, hostname, user, identity_file):

153

self.hostname = hostname

154

self.user = user

155

self.identity_file = os.path.expanduser(identity_file)

156

self._client = None

157

158

def get_private_key(self):

159

"""Generate a private key object for our keyfile"""

160

fp = os.path.expanduser(self.identity_file)

161

return paramiko.RSAKey.from_private_key(open(fp))

162

163

def connect(self):

164

info('Waiting for SSH to come available: %s@%s\n' % (

165

self.user, self.hostname))

166

debug("Using private key file: %s" % self.identity_file)

167

168

private_key = self.get_private_key()

169

170

for count in range(10):

171

self._client = paramiko.SSHClient()

172

self._client.set_missing_host_key_policy(ec2test.AcceptAllPolicy())

173

174

try:

175

176

self._client.connect(

177

self.hostname,

178

username=self.user,

179

pkey=private_key,

180

allow_agent=False,

181

look_for_keys=False)

182

183

except (socket.error, paramiko.AuthenticationException), e:

184

log.warning('wait_for_connection: %r' % (e,))

185

if count < 9:

186

time.sleep(5)

187

info('retrying...')

188

else:

189

raise

190

else:

191

break

192

193

def exec_command(self, remote_command, check_return=True):

194

"""Execute a command on the remote server.

195

196

Raises an error if the command returns an exit status that is not

197

zero, unless the option `check_return=False' has been given.

198

"""

199

info('Executing command: %s@%s %s\n' % (self.user, self.hostname, remote_command))

200

201

session = self._client.get_transport().open_session()

202

session.exec_command(remote_command)

203

session.shutdown_write()

204

205

# TODO: change this to use the logging module

206

while True:

207

select.select([session], [], [], 0.5)

208

if session.recv_ready():

209

data = session.recv(4096)

210

if data:

211

sys.stdout.write(data)

212

sys.stdout.flush()

213

if session.recv_stderr_ready():

214

data = session.recv_stderr(4096)

215

if data:

216

sys.stderr.write(data)

217

sys.stderr.flush()

218

if session.exit_status_ready():

219

break

220

session.close()

221

222

# XXX: JonathanLange 2009-05-31: If the command is killed by a signal

223

# on the remote server, the SSH protocol does not send an exit_status,

224

# it instead sends a different message with the number of the signal

225

# that killed the process. AIUI, this code will fail confusingly if

226

# that happens.

227

exit_status = session.recv_exit_status()

228

if exit_status and check_return:

229

raise RuntimeError('Command failed: %s' % (remote_command,))

230

return exit_status

231

232

def copy_to_remote(self, local_filename, remote_filename):

233

cmd = [

234

'scp',

235

'-i', self.identity_file,

236

local_filename,

237

'%s@%s:%s' % (self.user, self.hostname, remote_filename)

238

]

239

info("Executing command: %s" % ' '.join(cmd))

240

subprocess.check_call(cmd)

241

242

def user_command(self):

243

"""Return a user-friendly ssh command-line string."""

244

return "ssh -i %s %s@%s" % (

245

self.identity_file,

246

self.user,

247

self.hostname)

248

249

250

class ImageBundler:

251

"""Bundle an EC2 image on a remote system."""

252

253

def __init__(self, private_key, cert, account_id, target_bucket,

254

access_key, secret_key, ssh):

255

self.private_key = os.path.expanduser(private_key)

256

self.cert = os.path.expanduser(cert)

257

self.account_id = account_id

258

self.target_bucket = target_bucket

259

self.access_key = access_key

260

self.secret_key = secret_key

261

self.ssh = ssh

262

263

# Use the instance /mnt directory by default, because it has a few

264

# hundred GB of free space to work with.

265

self._bundle_dir = os.path.join('/mnt', target_bucket)

266

267

def bundle_image(self):

268

self.configure_bundling_environment()

269

manifest = self._bundle_image()

270

self._upload_bundle(manifest)

271

self._register_image(manifest)

272

273

def remote_private_keypath(self):

274

# ALWAYS have these files in /mnt on the remote system. Otherwise

275

# they will get bundled along with the image.

276

return os.path.join('/mnt', os.path.basename(self.private_key))

277

278

def remote_certpath(self):

279

# ALWAYS have these files in /mnt on the remote system. Otherwise

280

# they will get bundled along with the image.

281

return os.path.join('/mnt', os.path.basename(self.cert))

282

283

def configure_bundling_environment(self):

284

"""Configure what we need on the instance for bundling the image."""

285

# Send our keypair to the remote environment so that it can be used

286

# to bundle the image.

287

local_cert = os.path.abspath(self.cert)

288

local_pkey = os.path.abspath(self.private_key)

289

290

# ALWAYS copy these files into /mnt on the remote system. Otherwise

291

# they will get bundled along with the image.

292

remote_cert = self.remote_certpath()

293

remote_pkey = self.remote_private_keypath()

294

295

# See if the files are present, and copy them over if they are not.

296

self._ensure_remote_file(remote_cert, local_cert)

297

self._ensure_remote_file(remote_pkey, local_pkey)

298

299

def _ensure_remote_file(self, remote_file, desired_file):

300

info("Checking for '%s' on the remote system" % remote_file)

301

test = 'test -f %s' % remote_file

302

exit_status = self.ssh.exec_command(test, check_return=False)

303

if bool(exit_status):

304

self.ssh.copy_to_remote(desired_file, remote_file)

305

306

def _bundle_image(self):

307

# Create the bundle in a subdirectory, to avoid spamming up /mnt.

308

self.ssh.exec_command(

309

'mkdir %s' % self._bundle_dir, check_return=False)

310

311

cmd = [

312

'ec2-bundle-vol',

313

'-d %s' % self._bundle_dir,

314

'-b', # Set batch-mode, which doesn't use prompts.

315

'-k %s' % self.remote_private_keypath(),

316

'-c %s' % self.remote_certpath(),

317

'-u %s' % self.account_id

318

]

319

320

self.ssh.exec_command(' '.join(cmd))

321

# Assume that the manifest is 'image.manifest.xml', since "image" is

322

# the default prefix.

323

manifest = os.path.join(self._bundle_dir, 'image.manifest.xml')

324

325

# Best check that the manifest actually exists though.

326

test = 'test -f %s' % manifest

327

exit_status = self.ssh.exec_command(test, check_return=False)

328

329

if bool(exit_status):

330

raise RuntimeError(

331

"Failed to write the image manifest file: %s" % manifest)

332

333

return manifest

334

335

def _upload_bundle(self, manifest):

336

cmd = [

337

'ec2-upload-bundle',

338

'-b %s' % self.target_bucket,

339

'-m %s' % manifest,

340

'-a %s' % self.access_key,

341

'-s %s' % self.secret_key

342

]

343

self.ssh.exec_command(' '.join(cmd))

344

345

def _register_image(self, manifest):

346

# This is invoked locally.

347

mfilename = os.path.basename(manifest)

348

manifest_path = os.path.join(self.target_bucket, mfilename)

349

350

env = os.environ.copy()

351

env['JAVA_HOME'] = '/usr/lib/jvm/default-java'

352

cmd = [

353

'ec2-register',

354

'--private-key=%s' % self.private_key,

355

'--cert=%s' % self.cert,

356

manifest_path

357

]

358

info("Executing command: %s" % ' '.join(cmd))

359

subprocess.check_call(cmd, env=env)

360

361

362

class XvfbSystemConfigurator:

363

"""Configure a remote operating system over SSH to use the xvfb server."""

364

365

def __init__(self, ssh):

366

self.ssh = ssh

367

368

def configure_system(self):

369

"""Configure the operating system with the needed packages, etc."""

370

do = self.ssh.exec_command

371

372

# Make sure we know about all the packages, and where they may be

373

# found.

374

do("apt-get -y update")

375

# Install the necessary packages

376

do("apt-get -y install xvfb firefox xfonts-base")

377

378

379

class CombinedConfigParser:

380

"""Store and reconcile options for both optparse and ConfigParser."""

381

382

def __init__(self, optparser, cfgparser):

383

self._optparser = optparser

384

self._cfgparser = cfgparser

385

386

# A list of required optparse options.

387

self.required_options = []

388

self.known_options = []

389

390

# Our parsed positional command-line arguments, as returned by

391

# optparse.OptionParser.parse_args()

392

self.args = None

393

394

# An optparse option.dest to 'cfg-key' mapping.

395

self._option_to_cfgkey = {}

396

397

self._parsed_cli_options = None

398

self._parsed_cfg_options = None

399

400

def __getattr__(self, name):

401

return self.get(name)

402

403

def add_option(self, *args, **kwds):

404

"""Wrap the OptionParser.add_option() method, and add our options."""

405

try:

406

# We can't pass unknown kwds to make_option, or it will barf.

407

is_required = kwds.pop('required')

408

except KeyError:

409

is_required = False

410

411

option = optparse.make_option(*args, **kwds)

412

self._optparser.add_option(option)

413

414

if is_required:

415

self.add_required_option(option)

416

417

self._add_option_to_cfg_mapping(option, args)

418

419

def add_required_option(self, option):

420

"""Add a required option.

421

422

Takes an optparse.Option object.

423

"""

424

self.required_options.append(option)

425

426

def _add_option_to_cfg_mapping(self, option, option_constructor_args):

427

# Convert the long options into .ini keys. Use the last long option

428

# given.

429

for switch in reversed(option_constructor_args):

430

if switch.startswith('--'):

431

# We found a '--foo' switch, so use it. Drop the '--',

432

# because the config file doesn't use the prefixes.

433

self._option_to_cfgkey[option.dest] = switch[2:]

434

435

def error(self, message):

436

"""Wrap optparse.OptionParser.error()."""

437

self._optparser.error(message)

438

439

def parse_config_file(self, filepath):

440

fp = os.path.expanduser(filepath)

441

442

if not os.path.exists(fp):

443

self.error("The config file '%s' does not exist!" % fp)

444

445

self._cfgparser.read(fp)

446

self._parsed_cfg_options = self._cfgparser.defaults()

447

448

num_opts = len(self._parsed_cfg_options)

449

debug("Loaded %d options from %s" % (num_opts, fp))

450

451

def parse_cli_args(self, argv):

452

"""Wrap optparse.OptionParser.parse_args()."""

453

options, args = self._optparser.parse_args(argv)

454

self._parsed_cli_options = options

455

self.args = args

456

return (options, args)

457

458

def verify_options(self):

459

"""Verify that all required options are there.

460

461

Raise an optparse.OptionParser.error() if something is missing.

462

463

Make sure you parsed the config file with parse_config_file() before

464

doing this.

465

"""

466

debug("Verifying options")

467

if not self._parsed_cfg_options:

468

debug("No config file options found")

469

470

for option in self.required_options:

471

# Check for a command-line option.

472

473

option_name = option.dest

474

475

if self.get(option_name) is None:

476

self._required_option_error(option)

477

else:

478

debug("Found required option: %s" % option_name)

479

480

def _required_option_error(self, option):

481

msg = "Required option '%s' was not given (-h for help)" % str(option)

482

self.error(msg)

483

484

def get(self, name, default=None):

485

"""Return the appropriate option, CLI first, CFG second."""

486

cli_name = name

487

cfg_name = self._option_to_cfgkey.get(name)

488

489

value = self._getoption(cli_name)

490

491

if value is None and cfg_name is not None:

492

# No command-line option was supplied, but we do have a config

493

# file entry with that name.

494

value = self._getcfg(cfg_name)

495

496

if value is None:

497

# No config file option was supplied either, so return the

498

# default.

499

return default

500

501

return value

502

503

def _getoption(self, key, default=None):

504

return getattr(self._parsed_cli_options, key, default)

505

506

def _getcfg(self, key, default=None):

507

return self._parsed_cfg_options.get(key, default)

508

509

510

def get_credentials():

511

"""Return an EC2Credentials object for accessing the webservice."""

512

# Get the AWS identifier and secret identifier.

513

return ec2test.EC2Credentials.load_from_file()

514

515

516

def parse_config_file(filepath):

517

config = ConfigParser.ConfigParser()

518

config.read(filepath)

519

return config

520

521

522

def parse_options(argv):

523

oparser = optparse.OptionParser(usage)

524

cparser = ConfigParser.SafeConfigParser()

525

parser = CombinedConfigParser(oparser, cparser)

526

527

# What follows are "Required options" - these must be supplied from either

528

# the command-line, or from a config file.

529

parser.add_option(

530

'-k', '--key',

531

dest="keypair_name",

532

required=True,

533

help="The name of the AWS key pair to use for launching instances.")

534

535

parser.add_option(

536

'-K', '--private-key',

537

dest="private_key",

538

required=True,

539

help="The X.509 private keyfile that will be used to sign the new "

540

"image.")

541

542

parser.add_option(

543

'-C', '--cert',

544

dest="cert",

545

required=True,

546

help="The X.509 certificate that will be used to bundle the new "

547

"image.")

548

549

parser.add_option(

550

'-i', '--identity-file',

551

dest='identity_file',

552

required=True,

553

help="The location of the RSA private key that SSH will use to "

554

"connect to the instance.")

555

556

parser.add_option(

557

'-b', '--bucket',

558

dest="bucket",

559

required=True,

560

help="The bucket that the image will be placed into.")

561

562

parser.add_option(

563

'-u', '--user-id',

564

dest="account_id",

565

required=True,

566

help="Your 12 digit AWS account ID")

567

568

parser.add_option(

569

'-a', '--access-key',

570

dest="access_key",

571

required=True,

572

help="Your AWS access key.")

573

574

parser.add_option(

575

'-s', '--secret-key',

576

dest="secret_key",

577

required=True,

578

help="Your AWS secret key.")

579

580

581

# Start our "Optional options."

582

parser.add_option(

583

'-v', '--verbose',

584

action='store_true',

585

dest='verbose',

586

default=False,

587

help="Turn on debug output.")

588

589

parser.add_option(

590

'-t', '--instance-type',

591

dest="instance_type",

592

default="m1.large",

593

help="The type of instance to be launched. Should be the same as "

594

"the base image's required type. [default: %default]")

595

596

parser.add_option(

597

'-c', '--config',

598

dest='config',

599

default="~/.ec2/ec2bundle.cfg",

600

help="Load script options from the supplied config file. (.ini "

601

"format, see the module docstring for details.) "

602

"[default: %default]")

603

604

parser.add_option(

605

'--keepalive',

606

action='store_true',

607

dest="keepalive",

608

default=False,

609

help="Don't shut down the instance when we are done building (or "

610

"erroring out).")

611

612

parser.add_option(

613

'--no-bundle',

614

action='store_true',

615

dest='no_bundle',

616

default=False,

617

help="Don't create a bundle, just start the server and configure the "

618

"environment.")

619

620

parser.add_option(

621

'--use-instance',

622

dest='running_instance',

623

help="Use the supplied EC2 instance ID, instead of starting our own "

624

"server. The instance will be left running.")

625

626

627

options, args = parser.parse_cli_args(argv)

628

629

# Do this ASAP

630

if options.verbose:

631

log.setLevel(logging.DEBUG)

632

633

if options.config:

634

parser.parse_config_file(options.config)

635

636

# Make sure all the required args are present. Will error-out if

637

# something is missing.

638

parser.verify_options()

639

640

if len(args) != 2:

641

parser.error("You must provide an AMI ID that can serve as the new "

642

"image's base.")

643

644

return parser

645

646

647

def main(argv):

648

config = parse_options(argv)

649

650

credentials = get_credentials()

651

account = credentials.connect(config.keypair_name)

652

653

# Save the flag so we can change it. This is how we enforce shutdown

654

# policies.

655

keepalive = config.keepalive

656

657

if config.running_instance:

658

# We want to keep the server alive if the user supplied their own

659

# instance. Killing it without their consent would be cruel.

660

keepalive = True

661

662

ssh_user_command = None

663

try:

664

try:

665

instance = None

666

if config.running_instance:

667

# Connect to an already running instance.

668

instance = Instance.from_running_instance(

669

account, config.running_instance)

670

else:

671

# Run an instance for our base image.

672

instance = Instance.from_image(

673

account, config.args[1], config.instance_type)

674

675

ssh = SSHConnector(

676

instance.hostname, 'root', config.identity_file)

677

ssh.connect()

678

ssh_user_command = ssh.user_command()

679

680

system_configurator = XvfbSystemConfigurator(ssh)

681

system_configurator.configure_system()

682

683

if not config.no_bundle:

684

bundler = ImageBundler(

685

config.private_key,

686

config.cert,

687

config.account_id,

688

config.bucket,

689

config.access_key,

690

config.secret_key,

691

ssh)

692

bundler.bundle_image()

693

694

except:

695

# Log the exception now so it doesn't interfere with or get eaten

696

# by the instance shutdown.

697

log.exception("Oops!")

698

finally:

699

if keepalive:

700

log.warning("instance %s is now running on its own" % instance.id)

701

if ssh_user_command:

702

info("You may now ssh into the instance using the following command:")

703

info(" $ %s" % ssh_user_command)

704

705

log.warning("Remember to shut the instance down when you are done!")

706

else:

707

instance.stop()

708

709

710

if __name__ == '__main__':

711

logging.basicConfig()

712

main(sys.argv)