#!/usr/bin/python
# vim: expandtab:tabstop=4:shiftwidth=4
''' script to get ecr info '''
# Reason: disable invalid-name because pylint does not like our naming convention
# pylint: disable=invalid-name
import time
import boto3
import sys
import argparse
def get_volume(ec2, instanceId):
  result = []
  instance = ec2.Instance(instanceId)
  volumes = instance.volumes.all()
  for volume in volumes:
    print("Volume attached to this instance is :" + volume.id)
    result.append(volume.id)
  return result
def take_snapByInstance(client, instanceId):
  response = client.create_snapshots(
  Description='string',
  InstanceSpecification={
    'InstanceId': instanceId,
    'ExcludeBootVolume': False
  },
  TagSpecifications=[
    {
      'ResourceType': 'snapshot',
      'Tags': [
        {
          'Key': 'orginName',
          'Value': 'patch backup'+ instanceId
        },
      ]
    },
  ],
  DryRun=False,
  CopyTagsFromSource='volume'
  )
  print("Creating new snapshots for instances:" + response['Snapshots'][0]['SnapshotId'])
  return response['Snapshots'][0]['SnapshotId']
def get_publicIp(ec2, instanceId):
  instance = ec2.Instance(instanceId)
  publicIp = instance.public_ip_address
  return publicIp
def take_screenshotOfProcess(public_ip):
  print("Please run this command on your local machine")
  print('ssh -t ' + public_ip + ' "sudo netstat -tnpl > disk.listen"')
  print('ssh -t ' + public_ip + ' "sudo ps auxf > disk.ps"')
def get_elbInfo(client_elb, ec2, instanceId):
  bals = client_elb.describe_load_balancers()
  for elb in bals['LoadBalancerDescriptions']:
    #print('ELB DNS Name : ' + elb['DNSName'])
    #check if the elb is the elb of instance
    if instanceId in elb['Instances']:
      print("found elb " + elb['DNSName'])
    else:
      pass
def remove_fromElb(client_elb, elb, instanceId):
  response = client_elb.deregister_instances_from_load_balancer(
    LoadBalancerName='elb',
    Instances=[
      {
        'InstanceId': instanceId
      },
    ]
  )
def add_backElb(client_elb, elb, instanceId):
  response = client.register_instances_with_load_balancer(
    LoadBalancerName= elb,
    Instances=[
      {
        'InstanceId': instanceId
      },
    ]
  )
def check_snapStatus(ec2, snaps):
  snapshot = ec2.Snapshot(snaps)
  snapshot.load()
  print(snapshot.state)
  return snapshot.state
def main(ec2, client, instanceId, client_elb):
  print("going to paching instanceid: " + instanceId)
  #get volumes
  volumes = get_volume(ec2, instanceId)
  #get public ip
  public_ip = get_publicIp(ec2, instanceId)
  #take snapshot
  snaps = take_snapByInstance(client, instanceId)
  #take screenshot of procss and port
  take_screenshotOfProcess(public_ip)
  #get elb info
  elb = False
  #elb = get_elbInfo(client_elb, ec2, instanceId)
  #remove from elb
  if elb:
    ans_remove = input("Are you sure to remove the instance from the elb now")
    if ans_remove == 'Yes':
    #remove from instance
      remove_fromElb(client_elb, elb, instanceId)
  #check snapshot status
  snapshotStatus = ''
  check_snapStatus(ec2, snaps)
  print("checking staus of snapshots")
  while True:
    snapshotStatus = check_snapStatus(ec2, snaps)
    print(snapshotStatus)
    if snapshotStatus == 'completed':
      break
    else:
      time.sleep(10)
    #paching
  paching_cmd = 'Your paching command'
  print(paching_cmd)
  #add to elb
  if elb:
    ans_add = input("please confirm the patching is over , input yes to continue")
    if ans_add == 'Yes':
      add_backElb(client_elb, elb, instanceId)
if __name__ == "__main__":
  ec2 = boto3.resource('ec2', region_name='us-east-1')
  client = boto3.client('ec2', region_name='us-east-1')
  client_elb = boto3.client('elb', region_name='us-east-1')
  main(ec2, client, 'i-abcasdfa111122', client_elb)