#!/usr/bin/python
#
# Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Simple traceroute to the Gcloud VM's using gcloud SDK python library.

This script will lookup VMs within a gcloud project that matches a regexp
and run traceroute to each of them. This simplifies obtaining this kind
of debug information.

Typical Usage:
  python traceroute.py "[regexp]"

  Run "python traceroute.py -h" for more options

Requirements:
- Python 2.7
- Requires gcloud tool installed with a default project set up.
  This will be the project that the script will use to look for VMs.
- Requires the Google API Client for Python
  (https://cloud.google.com/compute/docs/tutorials/python-guide)
- Requires traceroute and dig (optional) to be installed and in PATH

"""

from __future__ import absolute_import
from __future__ import print_function

import argparse
import re
import subprocess

from googleapiclient import discovery
from oauth2client.client import GoogleCredentials

###############################################
# Argument parsing
###############################################


def parse_args():
  """Parse the command line arguments.

  Returns:
    The parsed argparse instance
  """
  parser = argparse.ArgumentParser(description="Utility for easily debug VMs")
  parser.add_argument(
      "-p",
      "--print",
      action="store_true",
      help="List all the instances, instead of traceroute.")
  parser.add_argument(
      "--project",
      help="""Project in which to run the tracerouter.
      If none is specified, the default project will be used.""")
  parser.add_argument(
      "-d",
      "--dig",
      action="store_true",
      help="Include dig (DNS lookup) information.")
  parser.add_argument(
      "-r",
      "--reverse_traceroute",
      action="store_true",
      help="Include reverse traceroute (from VM to host).")
  parser.add_argument(
      "match_pattern",
      help="""Pattern to match against VM names within the project.
                      Can use regexp expressions to match.""")

  return parser.parse_args()


##############################################
# gcloud account / helpers
##############################################


def get_gcloud_api():
  """Obtains the google SDK api instance.

  Returns:
    The authenticated gcloud api instance
  """
  credentials = GoogleCredentials.get_application_default()
  compute_api = discovery.build("compute", "v1", credentials=credentials)
  return compute_api


def obtain_self_ip():
  """Query open DNS to obtain the public IP of the current host.

  Returns:
    The public ip string
  """
  ip = subprocess.check_output(
      ["dig", "myip.opendns.com", "@resolver1.opendns.com", "+short"])
  return ip.strip()


def obtain_default_project():
  """Obtains the default project from the gcloud configuration.

  Returns:
    The name of the default project currently defined in gcloud's config
  """
  default_project = subprocess.check_output(
      ["gcloud", "config", "get-value", "project"], stderr=subprocess.PIPE)
  default_project = default_project.strip()
  return default_project


def list_instances(compute_api, project, zone):
  """List the instances for a project/zone.

  Args: