#!/usr/bin/env python2

#--Filename----------------------------------------------------------#
#   bmu.py                                                           #
#--Info--------------------------------------------------------------#
#   BlackArch mirrorlist update script                               #
#   Updated: 27/05/2015                                              #
#   The following lines:                                             #
#       [blackarch]                                                  #
#       Include = /etc/pacman.d/blackarch-mirrorlist                 #
#   Must be present in pacman.conf                                   #
#   blackarch.sh can be used to setup your pacman.conf correctly     #
#   Designed for Arch Linux - https://archlinux.org/download/        #
#   and BlackArch           - http://www.blackarch.org/index.html    #
#--Author------------------------------------------------------------#
#   JKAD - https://jkad.github.io                                    #
#--Tested------------------------------------------------------------#
#   27/05/2015 - archlinux-2015.05.01-dual.iso                       #
#--Licence-----------------------------------------------------------#
#   MIT Licence:                                                     #
#   https://github.com/JKAD/Arch-Install-Scripts/blob/master/LICENSE #
#--------------------------------------------------------------------#

import os
import sys
import urllib2
import time
from HTMLParser import HTMLParser
from datetime import datetime

protocol = "://"

class ParseHTML(HTMLParser):
    def __init__(self):
        HTMLParser.__init__(self)
        self.results = []
        self.return_data = False
        self.href = False

    def handle_starttag(self, tag, attrs):
        if tag == "img":
            for name, value in attrs:
                if "flags" in value:
                    self.return_data = True
        if tag == "href":
            self.href = True

    def handle_data(self, data):
        if self.href:
            if protocol in data:
                self.results.append(data)
                self.href = False
            self.return_data = False
        if self.return_data:
            self.results.append(data)

def main():
    if not os.geteuid() == 0:
            sys.exit('bmu.py must be run as root')

    url = "http://blackarch.org/downloads.html"
    user_agent = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 \
            (KHTML, like Gecko) Chrome/43.0.2357.81 Safari/537.36"
    FILE = '/etc/pacman.d/blackarch-mirrorlist'
    headers = {}
    headers["User-Agent"] = user_agent

    request = urllib2.Request(url, headers=headers)
    response = urllib2.urlopen(request)
    page = response.read()

    parser = ParseHTML()
    parser.feed(page)
    parser_results = parser.results
    ts = time.time()
    timestamp = "# Updated: %s\n\n" \
            % datetime.fromtimestamp(ts).strftime('%d-%m-%Y %H:%M:%S')
    message = "# Uncomment the mirrors you wish to use\n"
    default = "Server = http://mirror.team-cymru.org/blackarch/$repo/os/$arch\n"
    with open(FILE, 'w') as overwrite:
        overwrite.write(timestamp + message + default)
    for i in parser_results:
        if i.strip() != "":
            if protocol in i:
                hash_sym = "#Server = "
                ending =  "$repo/os/$arch\n"
            else:
                hash_sym = "\n#"
                ending = "\n"
            with open(FILE, 'a') as out_file:
                out_file.write(hash_sym + i.strip('\n\r\t ') + ending)
    print "Mirrorlist updated"
    print "Uncomment your prefered mirrors in %s and run pacman -Syy" % FILE
if __name__ == '__main__':
    main()
