0ctf geo newbie

0ctf_geo.py

import telnetlib
import pycountry
from geopy import GoogleV3
import re
import sys
import shelve

host = '202.112.26.111'
port = 29995
#google_geocoding_key = 'wat'

class Telnet(telnetlib.Telnet):
    # inherit from Telnetlib and add new method
    def __init__(self,host,port):
        telnetlib.Telnet.__init__(self,host,port)
    # make easier when you want to send raw data to server
    def writeRawData(self,data):
        return self.get_socket().send(data)

    def recvRawData(self,size):
        return self.get_socket().recv(size)
        
def alpha_2_by_location(place_name, country=False):
    # just some buggy code, carry on
    if (place_name == 'Holy See'): place_name = 'Vatican'
    if (place_name == 'Macedonia (the former Yugoslav Republic of)'): place_name = 'Macedonia'
    # lookup shelf first
    gapi = shelve.open('googly_cache', writeback=True)
    try:
        wat = place_name.encode('base64')
    except UnicodeEncodeError:
        wat = u' '.join(place_name).encode('utf-8').strip().encode('base64')
    if (wat in gapi):
        print('[*] Found in shelf')
        loc = gapi[wat]
    else:
        print('[*] Request from GGAPI')
        loc = geolocator.geocode(place_name).raw
        gapi[wat] = loc
        gapi.sync()
    gapi.close()
    # try to get the country's name
    try:
        for comp in loc['address_components']:
            if 'country' in comp['types']:
                country_name = comp['long_name']
                short_name = comp['short_name']
                break
        print('[*] Country is %s, shortname = %s' % (country_name, short_name))
    except:
        country_name = place_name
        print('[*] Geolocator failed, use %s as query' % country_name)
    # double check
    if country:
        if (not ((place_name in country_name) or (country_name in place_name))):
            country_name = place_name
    while True:
        try:
            print('[*] Trying to lookup %s' % country_name)
            country = pycountry.countries.get(name=country_name)
            print('[*] Found')
            res = country.alpha2
            break
        except KeyError:
            # chop off the last word
            chop = country_name.rsplit(' ', 1)
            if (chop[0] == country_name):
                # use short_name then
                res = short_name
                break
            country_name = chop[0]
            pass
    # wat
    if (place_name == "Mount Olympus"): res = 'GR'
    if (place_name == "Hyde Park"): res = 'GB'
    if (place_name == "Volga"): res = 'RU'
    if (place_name == "Alexandria"): res = 'EG'
    if (place_name == "Buenos Aires"): res = 'AR'
    if (place_name == "Sentosa"): res = 'SG'
    if (place_name == "Rickshaw capital of the world"): res = 'BD' # Dhaka
    if (place_name == "Cordova"): res = 'US'
    if (place_name == "Rajang"): res = 'MY'
    if (place_name == "Sanya"): res = 'CN'
    if (place_name == "Sydney"): res = 'AU'
    if (place_name == "Melboume"): res = 'AU'
    if (place_name == "Vancouver"): res = 'CA'
    if (place_name == "Lego"): res = 'SO' # tell me where it is :<
    if (place_name == "Naples"): res = 'IT'
    if (place_name == "Korea (Republic of)"): res = 'KR'
    if (place_name == "Virgin Islands (British)"): res = 'VG'
    if (place_name == "Korea (Democratic People's Republic of)"): res = 'KP'
    if (place_name == "Micronesia (Federated States of)"): res = 'FM'
    print(res)
    tel.writeRawData(res + '\n')

def read_place(country):
    s = tel.read_until(':')
    print(s)
    place_name = re.findall(r'---\n(.*):', s)[0]
    place_name = unicode(place_name, 'utf8')
    alpha_2_by_location(place_name, country)
    
def read_river():
    s = tel.read_until(':')
    print(s)
    if ('Andes' in s):
        res = 'VE\nCO\nEC\nPE\nAR\nBO\nCL'
    elif ('Mississippi River' in s):
        res = 'US'
    elif ('Greater Caucasus' in s):
        res = 'AZ\nGE\nRU'
    elif ('Nile' in s):
        res = 'TZ\nUG\nRW\nBI\nCD\nKE\nET\nER\nSS\nSD\nEG'
    elif ('Himalaya' in s):
        res = 'AF\nCN\nIN\nPK\nMM\nNP\nBT'
    elif ('Alps' in s):
        res = 'AT\nFR\nSI\nLI\nIT\nDE\nCH'
    elif ('Congo' in s):
        res = 'AO\nBI\nCM\nCF\nCG\nCD\nGA\nRW\nTZ\nZM'
    elif ('Danube' in s):
        res = 'DE\nAT\nSK\nHU\nHR\nRS\nBG\nRO\nMD\nUA'
    elif ('Amazon' in s):
        res = 'BR\nCO\nPE'
    elif ('Mekong' in s):
        res = 'TH\nKH\nCN\nLA\nMM\nVN'
    elif ('Rhine' in s):
        res = 'DE\nFR\nLI\nNL\nCH\nAT'
    elif ('Apennine' in s):
        res = 'IT\nSM'
    elif ('Parana' in s):
        res = 'AR\nBR\nPY'
    elif ('Rocky' in s):
        res = 'US\nCA'
    else:
        tel.interact()
    tel.writeRawData(res + '\n')
    print(res)
    for i in xrange(res.count('\n')): print tel.read_until(':'),

# init connection
tel = Telnet(host, port)
#geolocator = GoogleV3(api_key=google_geocoding_key)
geolocator = GoogleV3()
# Level 0 - Alpha-2 code by country name
for i in xrange(20):
    read_place(True)
# Level 1 - Alpha-2 code with places' name
for i in xrange(50):
    read_place(False)
# Level 2 - Alpha-2 code by river or mountain's name
for i in xrange(20):
    read_river()
# Interact
tel.interact()