<commit_before># Copyright 2011 Gilt Groupe, INC
#
# 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.
"""
various dns operations
"""

# import some useful stuff
import re
import os
import sys
import time
import shutil
import difflib
import datetime
import mothership.kv
import mothership.validate
import mothership.network_mapper
from mothership.mothership_models import *

class DNSError(Exception):
    pass

def generate_dns_header(cfg, origin, fqdn, realm, site_id, domain):
    """
    Returns zone header for fqdn
    """
    contact = cfg.contact
    serial = int(time.time())
    # try to get a realm-specific master
    nsmaster = mothership.kv.select(cfg, realm+'.'+site_id, key='nsmaster')
    if not nsmaster:
        # if we can't get a realm-specific master, look for a global one
        nsmaster = mothership.kv.select(cfg, None, key='nsmaster')
        if not nsmaster:
            raise DNSError("No nsmaster KV entry for \"%s.%s\", exiting" % (realm, site_id))
    nsmaster = nsmaster.value
    if '@' in contact: contact = contact.replace('@','.')
    if origin:
        header = '$ORIGIN %s.' % fqdn
    else:
        header = ''
    header += """
$TTL %s
@                       IN      SOA     %s. %s. (
                                        %-10d   ; Serial
                                        %-10d   ; Refresh
                                        %-10d   ; Retry
                                        %-10d   ; Expire
                                        %-10d   ; TTL
                                        )

""" % (cfg.zonettl, nsmaster, contact, serial, cfg.dns_refresh,
        cfg.dns_retry, cfg.dns_expire, cfg.zonettl)
    for rec in ['ns', 'mx']:
        header += generate_root_records(cfg, realm+'.'+site_id, domain, rec)
    return header

def generate_root_records(cfg, unqdn, domain, key):
    rec = mothership.kv.collect(cfg, unqdn, key)
    if not rec:
        raise DNSError('''
No %s defined for %s or %s, use:
    ship kv -a %s %s=<ns1>[,<nsN>]
to configure one or more %s
''' % (key, unqdn, domain, unqdn, key, key))
    for r in rec:
        slist = r.value.split(',')
        if r.hostname == unqdn:
            continue
    records = ''
    for server in slist:
        records += '%-20s\tIN\t%-8s%-16s\n' % ('', key.upper(), server)
    return records

def generate_dns_arecords(cfg, realm, site_id, domain):
    """
    Retrieves server list from mothership to create A records
    """
    alist = ''
    for n in cfg.dbsess.query(Network).\
        filter(Network.site_id==site_id).\
        filter(Network.realm==realm).all():
            if n.server_id:
                s = cfg.dbsess.query(Server).\
                    filter(Server.id==n.server_id).first()
                if n.ip and s.hostname:
                    alist += '%-20s\tIN\t%-8s%-16s\n' % (s.hostname, 'A', n.ip)
    return alist

def generate_dns_arpa(cfg, cidr, fqdn, realm, site_id, domain):
    """
    Retrieves server list from mothership to create arpa records
    """
    net = mothership.network_mapper.get_network(cidr)
    num = 0
    while re.search('\.0+$', net):
        net = re.sub('\.0+$', '', net)
        num += 1
    alist = ''
    for n in cfg.dbsess.query(Network).\
        filter(Network.site_id==site_id).\
        filter(Network.realm==realm).\
        order_by(Network.ip).all():
        if n.server_id:
            s = cfg.dbsess.query(Server).\
                    filter(Server.id==n.server_id).first()
            if mothership.network_mapper.within(n.ip, cidr) and n.ip and s.hostname:
                alist += '%-20s\tIN\t%-8s%s.%s.\n' % (
                    '.'.join(reversed(n.ip.split('.')[-num:])),
                    'PTR', s.hostname, fqdn)
    return net, alist

def generate_dns_addendum(cfg, realm, site_id, domain):
    """
    Turns the dns_addendum table into usable records
    """
    alist = ''
    for dns in cfg.dbsess.query(DnsAddendum).\
        filter(DnsAddendum.site_id==site_id).\
        filter(DnsAddendum.realm==realm).\
        order_by(DnsAddendum.record_type, DnsAddendum.host).all():
        if dns and dns.target:
            target = dns.target
        else:
            raise DNSError("Addendum record does not exist: %s" % realm+"."+site_id+"."+domain)
        if not re.search('\d+',dns.target.split('.')[-1]):
            if not target.endswith('.'):
                target += '.'
        alist += '%-20s\tIN\t%-8s%-16s\n' % (dns.host, dns.record_type, target)
    return alist

def generate_dns_output(cfg, domain, opts):
    """
    Creates DNS zonefiles
    """
    tmpdir = cfg.dns_tmpdir
    reload = False
    zones = []
    if opts.system:
        opts.outdir = tmpdir + cfg.zonedir
    if opts.outdir:
        if not os.path.exists(opts.outdir):
            os.makedirs(opts.outdir)
    if opts.all:
        print 'Generating ALL mothership DNS files'
        for site_id in cfg.site_ids:
            for realm in cfg.realms:
                fqn = mothership.validate.v_get_fqn(cfg, realm+'.'+site_id)
                zones.append(generate_dns_forward(cfg, fqn, opts))
                revzones = generate_dns_reverse(cfg, fqn, opts)
                if revzones:
                    for revzone in revzones:
                        zones.append(revzone)
                else:
                    print "No reverse zones created for %s" % (realm+'.'+site_id)
    else:
        zones.append(generate_dns_forward(cfg, domain, opts))
        revzones = generate_dns_reverse(cfg, fqn, opts)
        if revzones:
            for revzone in revzones:
                zones.append(revzone)
        else:
            print "No reverse zones created for %s" % (realm+'.'+site_id)
    # Remove empty zones
    zones = [ z for z in zones if z ]
    print zones
    validated = validate_zone_config(cfg, tmpdir, zones)
    reload = reload or validated
    if opts.system:
        validated = validate_zone_files(tmpdir, zones)
        if reload or validated:
            try:
                print 'Reloading named...',
                os.system('/sbin/service named reload')
                print 'done'
            except:
                raise DNSError('Error reloading named!')

def validate_zone_files(prefix, tmpzones):
    to_reload = False
    for tempzone in tmpzones:
        livezone = re.sub('^'+prefix, '', tempzone)
        compared = compare_files(livezone, tempzone)
        to_reload = to_reload or compared
    return to_reload

def validate_zone_config(cfg, prefix, tmpzones):
    tempconf = prefix + cfg.zonecfg
    zoneconf = ''
    for zone in tmpzones:
        name = os.path.basename(zone)
        zoneconf += create_zone_block(name)
    print 'Creating temporary zone config: %s' % tempconf
    if not os.path.exists(os.path.dirname(tempconf)):
        os.makedirs(os.path.dirname(tempconf))
    f = open(tempconf, 'w')
    f.write(zoneconf)
    f.close()
    return compare_files(cfg.zonecfg, tempconf)

def create_zone_block(name):
    return 'zone "%s" in {\n\ttype master;\n\tfile "%s";\n};\n\n' \
        % (name, name)

def confirm_change(prefix, zone):
    ans = raw_input('\nTo approve this change, type "%s_%s": ' % (prefix, zone))
    if ans != '%s_%s' % (prefix, zone):
        raise DNSError('Rollout of "%s" zone file aborted.' % (zone))
        print 'Rollout of "%s" zone file aborted.' % (zone)
        return False
    else:
        return True

def rollout_changes(oldfile, newfile):
    if os.path.exists(oldfile):
        backup = oldfile + '.bak'
        if os.path.exists(backup):
            print 'Removing old backup: %s' % backup
            os.remove(backup)
        stamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
        print 'Making snapshot: %s.%s' % (newfile, stamp)
        shutil.copy(oldfile, newfile+'.'+stamp)
        print 'Making backup: %s' % backup
        shutil.move(oldfile, backup)
    print 'Overwriting %s with %s' % (oldfile, newfile)
    if not os.path.exists(os.path.dirname(oldfile)):
        os.makedirs(os.path.dirname(oldfile))
    shutil.copy(newfile, oldfile)

def compare_files(oldfile, newfile):
    reload = False
    print 'Comparing %s %s' % (oldfile, newfile)
    # Ignore extra spaces for comparison
    olddata = []
    if os.path.exists(oldfile):
        olddata = sorted([ re.sub('\s+',' ',x)
            for x in open(oldfile).readlines() ])
    newdata = sorted([ re.sub('\s+',' ',x)
        for x in open(newfile).readlines() ])
    diff = difflib.Differ().compare(olddata, newdata)
    changes = False
    for d in diff:
        # exclude comments, blank/common lines, Serial
        if re.match('[-+]\s+(;|$)', d) \
            or not re.match('[-+]', d) \
            or 'Serial' in d:
            continue
        if not changes:
            print '-'*60
        changes = True
        print d
    if not changes:
        print 'Skipping %s, no changes discovered' % oldfile
        return reload
    else:
        print '-'*60
    if changes:
        if confirm_change('dns', os.path.basename(oldfile)):
            rollout_changes(oldfile, newfile)
            reload = True
    return reload

def generate_dns_forward(cfg, domain, opts):
    """
    Creates the forward zonefile for the specified domain
    """
    fqn = mothership.validate.v_get_fqn(cfg, domain)
    realm, site_id, domain = mothership.validate.v_split_fqn(fqn)
    forward = generate_dns_header(cfg, True, fqn, realm, site_id, domain)
    forward += generate_dns_arecords(cfg, realm, site_id, domain)
    forward += generate_dns_addendum(cfg, realm, site_id, domain)
    f = sys.stdout
    if opts.outdir:
        zone = '%s/%s' % (opts.outdir, fqn)
        print 'Writing DNS forward zone for %s to %s' % (fqn, zone)
        f = open(zone, 'w')
    else:
        print '\n' + '-'*60 + '\nDNS forward zone for %s:\n' % fqn  + '-'*60
    f.write(forward)
    if opts.outdir:
        f.close()
    if zone:
        return zone
    else:
        return None

def generate_dns_reverse(cfg, domain, opts):
    """
    Creates the reverse zonefiles for the specified domain
    """
    zones = []
    fqn = mothership.validate.v_get_fqn(cfg, domain)
    realm, site_id, domain = mothership.validate.v_split_fqn(fqn)
    netblocks = mothership.network_mapper.remap(cfg, 'cidr', dom='.'+fqn)
    if not netblocks:
        print "Skipping reverse zone for %s, undefined in mothership.yaml" % (realm+'.'+site_id)
        return False
    for cidr in netblocks:
        net, rev = generate_dns_arpa(cfg, cidr, fqn, realm, site_id, domain)
        net = '%s.in-addr.arpa' % '.'.join(reversed(net.split('.')))
        reverse = generate_dns_header(cfg, False, fqn, realm, site_id, domain)
        reverse += rev
        if opts.outdir:
            zone = '%s/%s' % (opts.outdir, net)
            print 'Writing DNS reverse zone for %s to %s' % (net, zone)
            f = open(zone, 'w')
            f.write(reverse)
            f.close()
            zones.append(zone)
        else:
            print '\n' + '-'*60 + '\nDNS reverse zone for %s:\n' % net  + '-'*60
            print reverse
    if zones:
        return zones
    else:
        return None

def update_table_dnsaddendum(cfg, info, delete=False):
    """
    updates the dns_addendum table
    can be used to add or delete dns records
    """
    if not delete:
        data = cfg.dbsess.query(DnsAddendum).\
            filter(DnsAddendum.site_id==info['site_id']).\
            filter(DnsAddendum.realm==info['realm']).\
            filter(DnsAddendum.target==info['target']).\
            filter(DnsAddendum.record_type==info['record_type']).\
            filter(DnsAddendum.host==info['host']).all()
        if data:
            print 'The following already exist in the database'
            for d in data:
                print '%d) %s IN %s %s' % (d.id, d.host, d.record_type, d.target)
            print 'Please delete first before adding again'
            return
        else:
            if len(info['target'].split('.')) < 4 and \
                (info['record_type'] == 'A' or info['record_type'] == 'CNAME'):
                print 'Target should end with ip or fqdn for A and CNAME types, aborting'
                return
            data = DnsAddendum(info['host'], info['record_type'],
                info['realm'], info['site_id'], info['target'])
            cfg.dbsess.add(data)
    else:
        idlist = []
        print 'For %s.%s, we found:' % (info['realm'], info['site_id'])
        for d in cfg.dbsess.query(DnsAddendum).\
            filter(DnsAddendum.site_id==info['site_id']).\
            filter(DnsAddendum.realm==info['realm']).\
            filter(DnsAddendum.host==info['host']).all():
            idlist.append(d.id)
            print '%d) %s IN %s %s' % (d.id, d.host, d.record_type, d.target)
        if len(idlist) == 0:
            print '    no DNS entries for %s.%s.%s' % (info['host'],
                info['realm'], info['site_id'])
        else:
            ans = raw_input('Which ID do you wish to delete? ')
            if not ans or int(ans) not in idlist:
                print 'Delete DNS entry aborted.'
                return
            else:
                print 'Deleting DNS entry %d' % int(ans)
                data = cfg.dbsess.query(DnsAddendum).\
                    filter(DnsAddendum.id==int(ans)).one()
                cfg.dbsess.delete(data)
    cfg.dbsess.commit()
<commit_msg>sprucing up forward generation code<commit_after># Copyright 2011 Gilt Groupe, INC
#
# 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.
"""
various dns operations
"""

# import some useful stuff
import re
import os
import sys
import time
import shutil
import difflib
import datetime
import mothership.kv
import mothership.validate
import mothership.network_mapper
from mothership.mothership_models import *

class DNSError(Exception):
    pass

def generate_dns_header(cfg, origin, fqdn, realm, site_id, domain):
    """
    Returns zone header for fqdn
    """
    contact = cfg.contact
    serial = int(time.time())
    # try to get a realm-specific master
    nsmaster = mothership.kv.select(cfg, realm+'.'+site_id, key='nsmaster')
    if not nsmaster:
        # if we can't get a realm-specific master, look for a global one
        nsmaster = mothership.kv.select(cfg, None, key='nsmaster')
        if not nsmaster:
            raise DNSError("No nsmaster KV entry for \"%s.%s\", exiting" % (realm, site_id))
    nsmaster = nsmaster.value
    if '@' in contact: contact = contact.replace('@','.')
    if origin:
        header = '$ORIGIN %s.' % fqdn
    else:
        header = ''
    header += """
$TTL %s
@                       IN      SOA     %s. %s. (
                                        %-10d   ; Serial
                                        %-10d   ; Refresh
                                        %-10d   ; Retry
                                        %-10d   ; Expire
                                        %-10d   ; TTL
                                        )

""" % (cfg.zonettl, nsmaster, contact, serial, cfg.dns_refresh,
        cfg.dns_retry, cfg.dns_expire, cfg.zonettl)
    for rec in ['ns', 'mx']:
        header += generate_root_records(cfg, realm+'.'+site_id, domain, rec)
    return header

def generate_root_records(cfg, unqdn, domain, key):
    rec = mothership.kv.collect(cfg, unqdn, key)
    if not rec:
        raise DNSError('''
No %s defined for %s or %s, use:
    ship kv -a %s %s=<ns1>[,<nsN>]
to configure one or more %s
''' % (key, unqdn, domain, unqdn, key, key))
    for r in rec:
        slist = r.value.split(',')
        if r.hostname == unqdn:
            continue
    records = ''
    for server in slist:
        records += '%-20s\tIN\t%-8s%-16s\n' % ('', key.upper(), server)
    return records

def generate_dns_arecords(cfg, realm, site_id, domain):
    """
    Retrieves server list from mothership to create A records
    """
    alist = ''
    for n in cfg.dbsess.query(Network).\
        filter(Network.site_id==site_id).\
        filter(Network.realm==realm).all():
            if n.server_id:
                s = cfg.dbsess.query(Server).\
                    filter(Server.id==n.server_id).first()
                if n.ip and s.hostname:
                    alist += '%-20s\tIN\t%-8s%-16s\n' % (s.hostname, 'A', n.ip)
    return alist

def generate_dns_arpa(cfg, cidr, fqdn, realm, site_id, domain):
    """
    Retrieves server list from mothership to create arpa records
    """
    net = mothership.network_mapper.get_network(cidr)
    num = 0
    while re.search('\.0+$', net):
        net = re.sub('\.0+$', '', net)
        num += 1
    alist = ''
    for n in cfg.dbsess.query(Network).\
        filter(Network.site_id==site_id).\
        filter(Network.realm==realm).\
        order_by(Network.ip).all():
        if n.server_id:
            s = cfg.dbsess.query(Server).\
                    filter(Server.id==n.server_id).first()
            if mothership.network_mapper.within(n.ip, cidr) and n.ip and s.hostname:
                alist += '%-20s\tIN\t%-8s%s.%s.\n' % (
                    '.'.join(reversed(n.ip.split('.')[-num:])),
                    'PTR', s.hostname, fqdn)
    return net, alist

def generate_dns_addendum(cfg, realm, site_id, domain):
    """
    Turns the dns_addendum table into usable records
    """
    alist = ''
    for dns in cfg.dbsess.query(DnsAddendum).\
        filter(DnsAddendum.site_id==site_id).\
        filter(DnsAddendum.realm==realm).\
        order_by(DnsAddendum.record_type, DnsAddendum.host).all():
        if dns and dns.target:
            target = dns.target
        else:
            raise DNSError("Addendum record does not exist: %s" % realm+"."+site_id+"."+domain)
        if not re.search('\d+',dns.target.split('.')[-1]):
            if not target.endswith('.'):
                target += '.'
        alist += '%-20s\tIN\t%-8s%-16s\n' % (dns.host, dns.record_type, target)
    return alist

def generate_dns_output(cfg, domain, opts):
    """
    Creates DNS zonefiles
    """
    tmpdir = cfg.dns_tmpdir
    reload = False
    zones = []
    if opts.system:
        opts.outdir = tmpdir + cfg.zonedir
    if opts.outdir:
        if not os.path.exists(opts.outdir):
            os.makedirs(opts.outdir)
    if opts.all:
        print 'Generating ALL mothership DNS files'
        for site_id in cfg.site_ids:
            for realm in cfg.realms:
                fqn = mothership.validate.v_get_fqn(cfg, realm+'.'+site_id)
                zones.append(generate_dns_forward(cfg, fqn, opts))
                revzones = generate_dns_reverse(cfg, fqn, opts)
                if revzones:
                    for revzone in revzones:
                        zones.append(revzone)
                else:
                    print "No reverse zones created for %s" % (realm+'.'+site_id)
    else:
        zones.append(generate_dns_forward(cfg, domain, opts))
        revzones = generate_dns_reverse(cfg, fqn, opts)
        if revzones:
            for revzone in revzones:
                zones.append(revzone)
        else:
            print "No reverse zones created for %s" % (realm+'.'+site_id)
    # Remove empty zones
    zones = [ z for z in zones if z ]
    validated = validate_zone_config(cfg, tmpdir, zones)
    reload = reload or validated
    if opts.system:
        validated = validate_zone_files(tmpdir, zones)
        if reload or validated:
            try:
                print 'Reloading named...',
                os.system('/sbin/service named reload')
                print 'done'
            except:
                raise DNSError('Error reloading named!')

def validate_zone_files(prefix, tmpzones):
    to_reload = False
    for tempzone in tmpzones:
        livezone = re.sub('^'+prefix, '', tempzone)
        compared = compare_files(livezone, tempzone)
        to_reload = to_reload or compared
    return to_reload

def validate_zone_config(cfg, prefix, tmpzones):
    tempconf = prefix + cfg.zonecfg
    zoneconf = ''
    for zone in tmpzones:
        name = os.path.basename(zone)
        zoneconf += create_zone_block(name)
    print 'Creating temporary zone config: %s' % tempconf
    if not os.path.exists(os.path.dirname(tempconf)):
        os.makedirs(os.path.dirname(tempconf))
    f = open(tempconf, 'w')
    f.write(zoneconf)
    f.close()
    return compare_files(cfg.zonecfg, tempconf)

def create_zone_block(name):
    return 'zone "%s" in {\n\ttype master;\n\tfile "%s";\n};\n\n' \
        % (name, name)

def confirm_change(prefix, zone):
    ans = raw_input('\nTo approve this change, type "%s_%s": ' % (prefix, zone))
    if ans != '%s_%s' % (prefix, zone):
        raise DNSError('Rollout of "%s" zone file aborted.' % (zone))
        print 'Rollout of "%s" zone file aborted.' % (zone)
        return False
    else:
        return True

def rollout_changes(oldfile, newfile):
    if os.path.exists(oldfile):
        backup = oldfile + '.bak'
        if os.path.exists(backup):
            print 'Removing old backup: %s' % backup
            os.remove(backup)
        stamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
        print 'Making snapshot: %s.%s' % (newfile, stamp)
        shutil.copy(oldfile, newfile+'.'+stamp)
        print 'Making backup: %s' % backup
        shutil.move(oldfile, backup)
    print 'Overwriting %s with %s' % (oldfile, newfile)
    if not os.path.exists(os.path.dirname(oldfile)):
        os.makedirs(os.path.dirname(oldfile))
    shutil.copy(newfile, oldfile)

def compare_files(oldfile, newfile):
    reload = False
    print 'Comparing %s %s' % (oldfile, newfile)
    # Ignore extra spaces for comparison
    olddata = []
    if os.path.exists(oldfile):
        olddata = sorted([ re.sub('\s+',' ',x)
            for x in open(oldfile).readlines() ])
    newdata = sorted([ re.sub('\s+',' ',x)
        for x in open(newfile).readlines() ])
    diff = difflib.Differ().compare(olddata, newdata)
    changes = False
    for d in diff:
        # exclude comments, blank/common lines, Serial
        if re.match('[-+]\s+(;|$)', d) \
            or not re.match('[-+]', d) \
            or 'Serial' in d:
            continue
        if not changes:
            print '-'*60
        changes = True
        print d
    if not changes:
        print 'Skipping %s, no changes discovered' % oldfile
        return reload
    else:
        print '-'*60
    if changes:
        if confirm_change('dns', os.path.basename(oldfile)):
            rollout_changes(oldfile, newfile)
            reload = True
    return reload

def generate_dns_forward(cfg, domain, opts):
    """
    Creates the forward zonefile for the specified domain
    """
    fqn = mothership.validate.v_get_fqn(cfg, domain)
    realm, site_id, domain = mothership.validate.v_split_fqn(fqn)
    forward = generate_dns_header(cfg, True, fqn, realm, site_id, domain)
    forward += generate_dns_arecords(cfg, realm, site_id, domain)
    forward += generate_dns_addendum(cfg, realm, site_id, domain)
    if opts.outdir:
        zone = '%s/%s' % (opts.outdir, fqn)
        print 'Writing DNS forward zone for %s to %s' % (fqn, zone)
        f = open(zone, 'w')
        f.write(forward)
        f.close()
    else:
        print '\n' + '-'*60 + '\nDNS forward zone for %s:\n' % fqn  + '-'*60
        print zone
    if zone:
        return zone
    else:
        return None

def generate_dns_reverse(cfg, domain, opts):
    """
    Creates the reverse zonefiles for the specified domain
    """
    zones = []
    fqn = mothership.validate.v_get_fqn(cfg, domain)
    realm, site_id, domain = mothership.validate.v_split_fqn(fqn)
    netblocks = mothership.network_mapper.remap(cfg, 'cidr', dom='.'+fqn)
    if not netblocks:
        print "Skipping reverse zone for %s, undefined in mothership.yaml" % (realm+'.'+site_id)
        return False
    for cidr in netblocks:
        net, rev = generate_dns_arpa(cfg, cidr, fqn, realm, site_id, domain)
        net = '%s.in-addr.arpa' % '.'.join(reversed(net.split('.')))
        reverse = generate_dns_header(cfg, False, fqn, realm, site_id, domain)
        reverse += rev
        if opts.outdir:
            zone = '%s/%s' % (opts.outdir, net)
            print 'Writing DNS reverse zone for %s to %s' % (net, zone)
            f = open(zone, 'w')
            f.write(reverse)
            f.close()
            zones.append(zone)
        else:
            print '\n' + '-'*60 + '\nDNS reverse zone for %s:\n' % net  + '-'*60
            print reverse
    if zones:
        return zones
    else:
        return None

def update_table_dnsaddendum(cfg, info, delete=False):
    """
    updates the dns_addendum table
    can be used to add or delete dns records
    """
    if not delete:
        data = cfg.dbsess.query(DnsAddendum).\
            filter(DnsAddendum.site_id==info['site_id']).\
            filter(DnsAddendum.realm==info['realm']).\
            filter(DnsAddendum.target==info['target']).\
            filter(DnsAddendum.record_type==info['record_type']).\
            filter(DnsAddendum.host==info['host']).all()
        if data:
            print 'The following already exist in the database'
            for d in data:
                print '%d) %s IN %s %s' % (d.id, d.host, d.record_type, d.target)
            print 'Please delete first before adding again'
            return
        else:
            if len(info['target'].split('.')) < 4 and \
                (info['record_type'] == 'A' or info['record_type'] == 'CNAME'):
                print 'Target should end with ip or fqdn for A and CNAME types, aborting'
                return
            data = DnsAddendum(info['host'], info['record_type'],
                info['realm'], info['site_id'], info['target'])
            cfg.dbsess.add(data)
    else:
        idlist = []
        print 'For %s.%s, we found:' % (info['realm'], info['site_id'])
        for d in cfg.dbsess.query(DnsAddendum).\
            filter(DnsAddendum.site_id==info['site_id']).\
            filter(DnsAddendum.realm==info['realm']).\
            filter(DnsAddendum.host==info['host']).all():
            idlist.append(d.id)
            print '%d) %s IN %s %s' % (d.id, d.host, d.record_type, d.target)
        if len(idlist) == 0:
            print '    no DNS entries for %s.%s.%s' % (info['host'],
                info['realm'], info['site_id'])
        else:
            ans = raw_input('Which ID do you wish to delete? ')
            if not ans or int(ans) not in idlist:
                print 'Delete DNS entry aborted.'
                return
            else:
                print 'Deleting DNS entry %d' % int(ans)
                data = cfg.dbsess.query(DnsAddendum).\
                    filter(DnsAddendum.id==int(ans)).one()
                cfg.dbsess.delete(data)
    cfg.dbsess.commit()
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
emailfilter.py -- Email tickets to Trac.

A simple MTA filter to create Trac tickets from inbound emails.

Copyright 2005, Daniel Lundin <daniel@edgewall.com>
Copyright 2005, Edgewall Software

The scripts reads emails from stdin and inserts directly into a Trac database.
MIME headers are mapped as follows:

    * From: => Reporter
    * Subject: => Summary
    * Body => Description

How to use
----------
 * Set TRAC_ENV_PATH to the path of your project's Trac environment
 * Configure script as a mail (pipe) filter with your MTA
    typically, this involves adding a line like this to /etc/aliases:
       somename: |/path/to/email2trac.py
    Check your MTA's documentation for specifics.

Todo
----
  * Configure target database through env variable?
  * Handle/discard HTML parts
  * Attachment support
"""

TRAC_ENV_PATH = '/var/trac/test'

import email
import sys

from trac.env import Environment
from trac.ticket import Ticket

class TicketEmailParser(object):

    env = None

    def __init__(self, env):
        self.env = env

    def parse(self, fp):
        msg = email.message_from_file(fp)
        tkt = Ticket(self.env)
        tkt['status'] = 'new'
        tkt['reporter'] = msg['from']
        tkt['summary'] = msg['subject']
        for part in msg.walk():
            if part.get_content_type() == 'text/plain':
                tkt['description'] = part.get_payload(decode=1).strip()

        if tkt.values.get('description'):
            tkt.insert()

if __name__ == '__main__':
    env = Environment(TRAC_ENV_PATH, create=0)
    tktparser = TicketEmailParser(env)
    tktparser.parse(sys.stdin)
<commit_msg>0.12dev: Added a comment that the script is only a starting point. Thanks to Gisbert Amm for the patch.<commit_after>#!/usr/bin/python
"""
emailfilter.py -- Email tickets to Trac.

A simple MTA filter to create Trac tickets from inbound emails.

Copyright 2005, Daniel Lundin <daniel@edgewall.com>
Copyright 2005, Edgewall Software

Please note:

This is only a starting point. See
 * http://trac.edgewall.org/ticket/5327 and
 * http://trac-hacks.org/intertrac/EmailtoTracScript
for more complete and advanced examples.

The scripts reads emails from stdin and inserts directly into a Trac database.
MIME headers are mapped as follows:

    * From: => Reporter
    * Subject: => Summary
    * Body => Description

How to use
----------
 * Set TRAC_ENV_PATH to the path of your project's Trac environment
 * Configure script as a mail (pipe) filter with your MTA
    typically, this involves adding a line like this to /etc/aliases:
       somename: |/path/to/email2trac.py
    Check your MTA's documentation for specifics.

Todo
----
  * Configure target database through env variable?
  * Handle/discard HTML parts
  * Attachment support
"""

TRAC_ENV_PATH = '/var/trac/test'

import email
import sys

from trac.env import Environment
from trac.ticket import Ticket

class TicketEmailParser(object):

    env = None

    def __init__(self, env):
        self.env = env

    def parse(self, fp):
        msg = email.message_from_file(fp)
        tkt = Ticket(self.env)
        tkt['status'] = 'new'
        tkt['reporter'] = msg['from']
        tkt['summary'] = msg['subject']
        for part in msg.walk():
            if part.get_content_type() == 'text/plain':
                tkt['description'] = part.get_payload(decode=1).strip()

        if tkt.values.get('description'):
            tkt.insert()

if __name__ == '__main__':
    env = Environment(TRAC_ENV_PATH, create=0)
    tktparser = TicketEmailParser(env)
    tktparser.parse(sys.stdin)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Sublime Text 3 Plugin to invoke Yapf on a python file.
"""
import codecs
import os
import re
import subprocess
import sys
import tempfile
import configparser

import sublime, sublime_plugin

KEY = "pyyapf"

def save_style_to_tempfile(in_dict):
    """
    Take a dictionary of yapf style settings and return the file
    name of a tempfile containing the expected config formatted
    style settings
    """

    cfg = configparser.RawConfigParser()
    cfg.add_section('style')
    for key in in_dict:
        cfg.set('style', key, in_dict[key])

    fobj, filename = tempfile.mkstemp()
    cfg.write(os.fdopen(fobj, "w"))
    return filename

def is_python(view):
    return view.score_selector(0, 'source.python') > 0

# pylint: disable=W0232
class YapfSelectionCommand(sublime_plugin.TextCommand):
    """
    The "yapf_selection" command formats the current selection (or the entire
    document if the "use_entire_file_if_no_selection" option is enabled).
    """

    def is_enabled(self):
        return is_python(self.view)

    encoding = None
    debug = False

    def encode_selection(self, selection):
        try:
            encoded = self.view.substr(selection).encode(self.encoding)
        except UnicodeEncodeError as err:
            msg = "You may need to re-open this file with a different encoding. Current encoding is %r." % self.encoding
            self.error("UnicodeEncodeError: %s\n\n%s", err, msg)
            return

        self.indent = b""
        detected = False
        unindented = []
        for line in encoded.splitlines(keepends=True):
            if not detected:
                codeline = line.strip()
                if len(codeline) > 0:
                    self.indent, _, _ = line.partition(codeline)
                    detected = True
            unindented.append(line[len(self.indent):])
        unindented = b''.join(unindented)
        return unindented

    def replace_selection(self, edit, selection, output):
        indent = self.indent.decode(self.encoding)
        reindented = []
        for line in output.splitlines():
            reindented.append(indent + line + '\n')
        self.view.replace(edit, selection, ''.join(reindented))

    def run(self, edit):
        """
        primary action when the plugin is triggered
        """
        print("Formatting selection with Yapf")

        settings = sublime.load_settings("PyYapf.sublime-settings")

        self.encoding = self.view.encoding()

        if self.encoding == "Undefined":
            print('Encoding is not specified.')
            self.encoding = settings.get('default_encoding')

        print('Using encoding of %r' % self.encoding)

        self.debug = settings.get('debug')

        # there is always at least one region
        for region in self.view.sel():
            # determine selection to format
            if region.empty():
                if settings.get("use_entire_file_if_no_selection"):
                    selection = sublime.Region(0, self.view.size())
                else:
                    sublime.error_message('A selection is required')
                    continue
            else:
                selection = region

            # encode selection
            encoded_selection = self.encode_selection(selection)
            if not encoded_selection:
                continue

            # determine yapf command
            cmd = settings.get("yapf_command")
            assert cmd, "yapf_command not configured"
            cmd = os.path.expanduser(cmd)
            args = [cmd]

            # verify reformatted code
            args += ["--verify"]

            # override style?
            if settings.has('config'):
                custom_style = settings.get("config")
                style_filename = save_style_to_tempfile(custom_style)
                args += ["--style={0}".format(style_filename)]

                if self.debug:
                    print('Using custom style:')
                    with open(style_filename) as file_handle:
                        print(file_handle.read())
            else:
                style_filename = None

            # use directory of current file so that custom styles are found properly
            fname = self.view.file_name()
            cwd = os.path.dirname(fname) if fname else None

            # specify encoding in environment
            env = os.environ.copy()
            env['LANG'] = self.encoding

            # win32: hide console window
            if sys.platform in ('win32', 'cygwin'):
                startupinfo = subprocess.STARTUPINFO()
                startupinfo.dwFlags = subprocess.CREATE_NEW_CONSOLE | subprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = subprocess.SW_HIDE
            else:
                startupinfo = None

            # run yapf
            print('Running {0} in {1}'.format(args, cwd))
            if self.debug:
                print('Environment: {0}'.format(env))
            popen = subprocess.Popen(args,
                                     stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE,
                                     stdin=subprocess.PIPE,
                                     cwd=cwd,
                                     env=env,
                                     startupinfo=startupinfo)
            encoded_output, encoded_err = popen.communicate(encoded_selection)

            # handle errors (since yapf>=0.3, exit code 2 means changed, not error)
            if popen.returncode not in (0, 2):
                err = encoded_err.decode(self.encoding)
                print('Error:\n%s', err)

                # report error
                err_lines = err.splitlines()
                msg = err_lines[-1]
                if 'InternalError' in msg:
                    sublime.error_message(msg)
                else:
                    loc = err_lines[-4]
                    loc = loc[loc.find('line'):].capitalize()
                    sublime.error_message('%s (%s)' % (msg, loc))
            else:
                new_text = encoded_output.decode(self.encoding)
                self.replace_selection(edit, selection, new_text)

            if style_filename:
                os.unlink(style_filename)

        print('PyYapf Completed')

# pylint: disable=W0232
class YapfDocumentCommand(sublime_plugin.TextCommand):
    """
    The "yapf_document" command formats the current document.
    """

    def is_enabled(self):
        return is_python(self.view)

    def run(self, edit):
        # XXX
        self.view.run_command('yapf_selection')

class EventListener(sublime_plugin.EventListener):
    def on_pre_save(self, view):
        settings = sublime.load_settings("PyYapf.sublime-settings")
        if settings.get('on_save'):
            view.run_command('yapf_document')
<commit_msg>Refactored common code into Yapf class<commit_after># -*- coding: utf-8 -*-
"""
Sublime Text 3 Plugin to invoke Yapf on a python file.
"""
import codecs
import os
import re
import subprocess
import sys
import tempfile
import configparser

import sublime, sublime_plugin

KEY = "pyyapf"

def save_style_to_tempfile(style):
    # build config object
    cfg = configparser.RawConfigParser()
    cfg.add_section('style')
    for key, value in style.items():
        cfg.set('style', key, value)

    # dump it to temporary file
    fobj, fname = tempfile.mkstemp()
    cfg.write(os.fdopen(fobj, "w"))
    return fname

def is_python(view):
    return view.score_selector(0, 'source.python') > 0

class Yapf:
    """
    This class wraps YAPF invocation, including encoding/decoding and error handling.
    """

    def __init__(self, view):
        self.settings = sublime.load_settings("PyYapf.sublime-settings")
        self.view = view

    encoding = None
    debug = False

    def encode_selection(self, selection):
        try:
            encoded = self.view.substr(selection).encode(self.encoding)
        except UnicodeEncodeError as err:
            msg = "You may need to re-open this file with a different encoding. Current encoding is %r." % self.encoding
            self.error("UnicodeEncodeError: %s\n\n%s", err, msg)
            return

        self.indent = b""
        detected = False
        unindented = []
        for line in encoded.splitlines(keepends=True):
            if not detected:
                codeline = line.strip()
                if len(codeline) > 0:
                    self.indent, _, _ = line.partition(codeline)
                    detected = True
            unindented.append(line[len(self.indent):])
        unindented = b''.join(unindented)
        return unindented

    def replace_selection(self, edit, selection, output):
        indent = self.indent.decode(self.encoding)
        reindented = []
        for line in output.splitlines():
            reindented.append(indent + line + '\n')
        self.view.replace(edit, selection, ''.join(reindented))

    def format(self, selection, edit):
        """
        primary action when the plugin is triggered
        """
        print("Formatting selection with Yapf")

        self.encoding = self.view.encoding()

        if self.encoding == "Undefined":
            print('Encoding is not specified.')
            self.encoding = self.settings.get('default_encoding')

        print('Using encoding of %r' % self.encoding)

        self.debug = self.settings.get('debug')

        # encode selection
        encoded_selection = self.encode_selection(selection)
        if not encoded_selection:
            return

        # determine yapf command
        cmd = self.settings.get("yapf_command")
        assert cmd, "yapf_command not configured"
        cmd = os.path.expanduser(cmd)
        args = [cmd]

        # verify reformatted code
        args += ["--verify"]

        # override style?
        if self.settings.has('config'):
            custom_style = self.settings.get("config")
            style_filename = save_style_to_tempfile(custom_style)
            args += ["--style={0}".format(style_filename)]

            if self.debug:
                print('Using custom style:')
                with open(style_filename) as file_handle:
                    print(file_handle.read())
        else:
            style_filename = None

        # use directory of current file so that custom styles are found properly
        fname = self.view.file_name()
        cwd = os.path.dirname(fname) if fname else None

        # specify encoding in environment
        env = os.environ.copy()
        env['LANG'] = self.encoding

        # win32: hide console window
        if sys.platform in ('win32', 'cygwin'):
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags = subprocess.CREATE_NEW_CONSOLE | subprocess.STARTF_USESHOWWINDOW
            startupinfo.wShowWindow = subprocess.SW_HIDE
        else:
            startupinfo = None

        # run yapf
        print('Running {0} in {1}'.format(args, cwd))
        if self.debug:
            print('Environment: {0}'.format(env))
        popen = subprocess.Popen(args,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 stdin=subprocess.PIPE,
                                 cwd=cwd,
                                 env=env,
                                 startupinfo=startupinfo)
        encoded_output, encoded_err = popen.communicate(encoded_selection)

        # handle errors (since yapf>=0.3, exit code 2 means changed, not error)
        if popen.returncode not in (0, 2):
            err = encoded_err.decode(self.encoding)
            print('Error:\n%s', err)

            # report error
            err_lines = err.splitlines()
            msg = err_lines[-1]
            if 'InternalError' in msg:
                sublime.error_message(msg)
            else:
                loc = err_lines[-4]
                loc = loc[loc.find('line'):].capitalize()
                sublime.error_message('%s (%s)' % (msg, loc))
        else:
            new_text = encoded_output.decode(self.encoding)
            self.replace_selection(edit, selection, new_text)

        if style_filename:
            os.unlink(style_filename)

        print('PyYapf Completed')

# pylint: disable=W0232
class YapfSelectionCommand(sublime_plugin.TextCommand):
    """
    The "yapf_selection" command formats the current selection (or the entire
    document if the "use_entire_file_if_no_selection" option is enabled).
    """

    def is_enabled(self):
        return is_python(self.view)

    def run(self, edit):
        yapf = Yapf(self.view)

        # there is always at least one region...
        for region in self.view.sel():
            # ...if it is empty then there is no selection
            if region.empty():
                if yapf.settings.get("use_entire_file_if_no_selection"):
                    selection = sublime.Region(0, self.view.size())
                else:
                    sublime.error_message('A selection is required')
                    break
            else:
                selection = region

            # format region
            yapf.format(selection, edit)

# pylint: disable=W0232
class YapfDocumentCommand(sublime_plugin.TextCommand):
    """
    The "yapf_document" command formats the current document.
    """

    def is_enabled(self):
        return is_python(self.view)

    def run(self, edit):
        selection = sublime.Region(0, self.view.size())
        Yapf(self.view).format(selection, edit)

class EventListener(sublime_plugin.EventListener):
    def on_pre_save(self, view):
        settings = sublime.load_settings("PyYapf.sublime-settings")
        if settings.get('on_save'):
            view.run_command('yapf_document')
<|endoftext|>


<commit_before>""" CLI logger """
import logging
import logging.handlers
import os
import sys

import requests

import coloredlogs

try:
    import colorama
except ImportError:
    # coloredlogs only installs colorama on Windows
    pass

def init_logger(log_requests=False):
    """ Initialize the logger """

    logger = logging.getLogger(__name__.split(".")[0])
    for handler in logger.handlers:  # pragma: no cover
        logger.removeHandler(handler)

    if os.name == "nt" and "colorama" in sys.modules:  # pragma: no cover
        colorama.init()

    formatter = coloredlogs.ColoredFormatter(fmt="%(asctime)s: %(message)s")
    handler = logging.StreamHandler()
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)
    logger.propagate = False

    if log_requests:
        requests.packages.urllib3.add_stderr_logger()

def get_gist_logger(repo_root):
    """Determines the appropriate filepath for logfile
    and name for the logger. Returns a logger with
    RotatingFileHandler attached."""
    if repo_root:
        logfile_path = f"{repo_root}/.cci/cci.log"
        # create .cci/ if it doesn't exist
        if not os.path.exists(os.path.dirname(logfile_path)):
            os.makedirs(os.path.dirnmae(logfile_path))
    else:
        logfile_path = "cci.log"
    return get_rot_file_logger("sys.stdout", logfile_path)

def get_rot_file_logger(name, path):
    """Returns a logger with a rotating file handler"""
    logger = logging.getLogger(name)

    handler = logging.handlers.RotatingFileHandler(path, backupCount=5)
    handler.doRollover()  # rollover existing log files
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)
    return logger
<commit_msg>use better syntax to create .cci dir<commit_after>""" CLI logger """
import logging
import logging.handlers
import os
import sys

import requests

import coloredlogs

try:
    import colorama
except ImportError:
    # coloredlogs only installs colorama on Windows
    pass

def init_logger(log_requests=False):
    """ Initialize the logger """

    logger = logging.getLogger(__name__.split(".")[0])
    for handler in logger.handlers:  # pragma: no cover
        logger.removeHandler(handler)

    if os.name == "nt" and "colorama" in sys.modules:  # pragma: no cover
        colorama.init()

    formatter = coloredlogs.ColoredFormatter(fmt="%(asctime)s: %(message)s")
    handler = logging.StreamHandler()
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)
    logger.propagate = False

    if log_requests:
        requests.packages.urllib3.add_stderr_logger()

def get_gist_logger(repo_root):
    """Determines the appropriate filepath for logfile
    and name for the logger. Returns a logger with
    RotatingFileHandler attached."""
    if repo_root:
        logfile_path = f"{repo_root}/.cci/cci.log"
        # create .cci/ if it doesn't exist
        os.makedirs(os.path.dirname(logfile_path), exist_ok=True)
    else:
        logfile_path = "cci.log"
    return get_rot_file_logger("sys.stdout", logfile_path)

def get_rot_file_logger(name, path):
    """Returns a logger with a rotating file handler"""
    logger = logging.getLogger(name)

    handler = logging.handlers.RotatingFileHandler(path, backupCount=5)
    handler.doRollover()  # rollover existing log files
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)
    return logger
<|endoftext|>


<commit_before>#!/usr/bin/env python
import json
import os
import shlex
import socket
import sys

from docker import Client
from docker import errors

from paasta_tools.marathon_tools import CONTAINER_PORT
from paasta_tools.marathon_tools import get_args
from paasta_tools.marathon_tools import get_cpus
from paasta_tools.marathon_tools import get_mem
from paasta_tools.marathon_tools import read_service_config
from paasta_tools.marathon_tools import read_service_namespace_config
from paasta_tools.paasta_cli.utils import figure_out_service_name
from paasta_tools.paasta_cli.utils import lazy_choices_completer
from paasta_tools.paasta_cli.utils import list_instances
from paasta_tools.paasta_cli.utils import list_services
from paasta_tools.utils import PaastaColors
from paasta_tools.utils import read_marathon_config

def pick_random_port():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(('localhost', 0))
    addr, port = s.getsockname()
    s.close()
    return port

def get_healthcheck(service, instance, random_port):
    smartstack_config = read_service_namespace_config(service, instance)
    mode = smartstack_config.get('mode', 'http')
    hostname = socket.getfqdn()

    uri = PaastaColors.cyan('%s://%s:%d\n' % (mode, hostname, random_port))

    return 'Mesos would have healthchecked your service via\n%s\n' % uri

def add_subparser(subparsers):
    list_parser = subparsers.add_parser(
        'test-run',
        description='Test run service Docker container',
        help='Test run service Docker container',
    )
    list_parser.add_argument(
        '-s', '--service',
        help='The name of the service you wish to inspect',
    ).completer = lazy_choices_completer(list_services)
    list_parser.add_argument(
        '-C', '--cmd',
        help=(
            'Run Docker container with particular command, '
            'for instance bash'
        ),
        required=False,
        default='',
    )
    list_parser.add_argument(
        '-i', '--instance',
        help='Simulate a docker run for a particular instance of the service, like "main" or "canary"',
        required=False,
        default='main',
    ).completer = lazy_choices_completer(list_instances)
    list_parser.add_argument(
        '-v', '--verbose',
        help='Show Docker commands output',
        action='store_true',
        required=False,
        default=True,
    )
    list_parser.add_argument(
        '-I', '--interactive',
        help='Run container in interactive mode',
        action='store_true',
        required=False,
        default=False,
    )

    list_parser.set_defaults(command=paasta_test_run)

def run_docker_container_interactive(service, instance, docker_hash, volumes, command, service_manifest):
    """
    Since docker-py has some issues with running a container with TTY attached we're just executing
    docker run command for interactive mode. For non-interactive mode we're using docker-py.
    """
    sys.stderr.write(PaastaColors.yellow(
        'Warning! You\'re running a container in interactive mode.\n'
        'This is not how Mesos runs containers. To run container exactly\n'
        'like Mesos does don\'t use the -I flag.\n\n'
    ))

    run_args = ['docker', 'run']

    for volume in volumes:
        run_args.append('--volume=%s' % volume)

    random_port = pick_random_port()

    run_args.append('--tty=true')
    run_args.append('--interactive=true')
    run_args.append('--memory=%s' % (str(get_mem(service_manifest)) + 'm'))
    run_args.append('--publish=%d:%d' % (random_port, CONTAINER_PORT))
    run_args.append('%s' % docker_hash)
    run_args.extend(command)

    sys.stdout.write('Running docker command:\n%s\n' % ' '.join(run_args))

    healthcheck_string = get_healthcheck(service, instance, random_port)
    sys.stdout.write(healthcheck_string)

    os.execlp('/usr/bin/docker', *run_args)

def run_docker_container_non_interactive(
    docker_client,
    service,
    instance,
    docker_hash,
    volumes,
    command,
    service_manifest
):
    """
    Using docker-py for non-interactive run of a container. In the end of function it stops the container
    and removes it.
    """
    sys.stderr.write(PaastaColors.yellow(
        'Warning! You\'re running a container in non-interactive mode.\n'
        'This is how Mesos runs containers. Some programs behave differently\n'
        'with no tty attach.\n\n'
    ))

    create_result = docker_client.create_container(
        image=docker_hash,
        command=command,
        tty=False,
        volumes=volumes,
        mem_limit=str(get_mem(service_manifest)) + 'm',
        cpu_shares=get_cpus(service_manifest),
        ports=[CONTAINER_PORT],
        stdin_open=False,
    )

    docker_client.start(create_result['Id'], port_bindings={CONTAINER_PORT: None})

    smartstack_config = read_service_namespace_config(service, instance)
    port = smartstack_config.get('proxy_port', 0)

    healthcheck_string = get_healthcheck(service, instance, port)
    sys.stdout.write(healthcheck_string)

    for line in docker_client.attach(create_result['Id'], stream=True, logs=True):
        sys.stdout.write(line)

    docker_client.stop(create_result['Id'])
    docker_client.remove_container(create_result['Id'])

def run_docker_container(docker_client, docker_hash, service, args):
    """
    Run Docker container by image hash with args set in command line.
    Function prints the output of run command in stdout.
    """
    marathon_config_raw = read_marathon_config()

    volumes = list()

    for volume in marathon_config_raw['docker_volumes']:
        volumes.append('%s:%s:%s' % (volume['hostPath'], volume['containerPath'], volume['mode'].lower()))

    marathon_config = dict()

    marathon_config['cluster'] = marathon_config_raw['cluster']
    marathon_config['volumes'] = volumes

    service_manifest = read_service_config(service, args.instance, marathon_config['cluster'])

    if args.cmd:
        command = shlex.split(args.cmd)
    else:
        command = get_args(service_manifest)

    if args.interactive:
        run_docker_container_interactive(service, args.instance, docker_hash, volumes, command, service_manifest)
    else:
        run_docker_container_non_interactive(
            docker_client,
            service,
            args.instance,
            docker_hash,
            volumes,
            command,
            service_manifest
        )

def build_docker_container(docker_client, args):
    """
    Build Docker container from Dockerfile in the current directory or
    specified in command line args. Resulting image hash.
    """
    result = ''

    dockerfile_path = os.getcwd()

    sys.stdout.write('Building container from Dockerfile in %s\n' % dockerfile_path)

    for line in docker_client.build(path=dockerfile_path, tag='latest'):
        line_dict = json.loads(line)

        stream_line = line_dict.get('stream')

        if args.verbose:
            sys.stdout.write(stream_line)

        if stream_line.startswith('Successfully built '):
            """Strip the beginning of a string and \n in the end."""
            result = stream_line[len('Successfully built '):]
            result = result[:len(result) - 1]

    return result

def paasta_test_run(args):
    service = figure_out_service_name(args)

    base_docker_url = os.environ.get('DOCKER_HOST', 'unix://var/run/docker.sock')

    docker_client = Client(base_url=base_docker_url)

    try:
        docker_hash = build_docker_container(docker_client, args)
    except errors.APIError as e:
        sys.stderr.write('Can\'t build Docker container. Error: %s\n' % str(e))
        sys.exit(1)

    try:
        run_docker_container(docker_client, docker_hash, service, args)
    except errors.APIError as e:
        sys.stderr.write('Can\'t run Docker container. Error: %s\n' % str(e))
        sys.exit(1)
<commit_msg>Some CR issues.<commit_after>#!/usr/bin/env python
import json
import os
import shlex
import socket
import sys

from docker import Client
from docker import errors

from paasta_tools.marathon_tools import CONTAINER_PORT
from paasta_tools.marathon_tools import get_args
from paasta_tools.marathon_tools import get_cpus
from paasta_tools.marathon_tools import get_mem
from paasta_tools.marathon_tools import read_service_config
from paasta_tools.marathon_tools import read_service_namespace_config
from paasta_tools.paasta_cli.utils import figure_out_service_name
from paasta_tools.paasta_cli.utils import lazy_choices_completer
from paasta_tools.paasta_cli.utils import list_instances
from paasta_tools.paasta_cli.utils import list_services
from paasta_tools.utils import PaastaColors
from paasta_tools.utils import read_marathon_config

def pick_random_port():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(('localhost', 0))
    addr, port = s.getsockname()
    s.close()
    return port

def get_healthcheck(service, instance, random_port):
    smartstack_config = read_service_namespace_config(service, instance)
    mode = smartstack_config.get('mode', 'http')
    hostname = socket.getfqdn()

    uri = PaastaColors.cyan('%s://%s:%d\n' % (mode, hostname, random_port))

    return 'Mesos would have healthchecked your service via\n%s\n' % uri

def add_subparser(subparsers):
    list_parser = subparsers.add_parser(
        'test-run',
        description='Test run service Docker container',
        help='Test run service Docker container',
    )
    list_parser.add_argument(
        '-s', '--service',
        help='The name of the service you wish to inspect',
    ).completer = lazy_choices_completer(list_services)
    list_parser.add_argument(
        '-C', '--cmd',
        help=(
            'Run Docker container with particular command, '
            'for instance bash'
        ),
        required=False,
        default='',
    )
    list_parser.add_argument(
        '-i', '--instance',
        help='Simulate a docker run for a particular instance of the service, like "main" or "canary"',
        required=False,
        default='main',
    ).completer = lazy_choices_completer(list_instances)
    list_parser.add_argument(
        '-v', '--verbose',
        help='Show Docker commands output',
        action='store_true',
        required=False,
        default=True,
    )
    list_parser.add_argument(
        '-I', '--interactive',
        help='Run container in interactive mode',
        action='store_true',
        required=False,
        default=False,
    )

    list_parser.set_defaults(command=paasta_test_run)

def run_docker_container_interactive(service, instance, docker_hash, volumes, command, service_manifest):
    """
    Since docker-py has some issues with running a container with TTY attached we're just executing
    docker run command for interactive mode. For non-interactive mode we're using docker-py.
    """
    sys.stderr.write(PaastaColors.yellow(
        'Warning! You\'re running a container in interactive mode.\n'
        'This is not how Mesos runs containers. To run container exactly\n'
        'like Mesos does don\'t use the -I flag.\n\n'
    ))

    run_args = ['docker', 'run']

    for volume in volumes:
        run_args.append('--volume=%s' % volume)

    random_port = pick_random_port()

    run_args.append('--tty=true')
    run_args.append('--interactive=true')
    run_args.append('--memory=%dm' % get_mem(service_manifest))
    run_args.append('--publish=%d:%d' % (random_port, CONTAINER_PORT))
    run_args.append('%s' % docker_hash)
    run_args.extend(command)

    sys.stdout.write('Running docker command:\n%s\n' % ' '.join(run_args))

    healthcheck_string = get_healthcheck(service, instance, random_port)
    sys.stdout.write(healthcheck_string)

    os.execlp('/usr/bin/docker', *run_args)

def run_docker_container_non_interactive(
    docker_client,
    service,
    instance,
    docker_hash,
    volumes,
    command,
    service_manifest
):
    """
    Using docker-py for non-interactive run of a container. In the end of function it stops the container
    and removes it.
    """
    sys.stderr.write(PaastaColors.yellow(
        'Warning! You\'re running a container in non-interactive mode.\n'
        'This is how Mesos runs containers. Some programs behave differently\n'
        'with no tty attach.\n\n'
    ))

    create_result = docker_client.create_container(
        image=docker_hash,
        command=command,
        tty=False,
        volumes=volumes,
        mem_limit='%dm' % get_mem(service_manifest),
        cpu_shares=get_cpus(service_manifest),
        ports=[CONTAINER_PORT],
        stdin_open=False,
    )

    docker_client.start(create_result['Id'], port_bindings={CONTAINER_PORT: None})

    smartstack_config = read_service_namespace_config(service, instance)
    port = smartstack_config.get('proxy_port', 0)

    healthcheck_string = get_healthcheck(service, instance, port)
    sys.stdout.write(healthcheck_string)

    for line in docker_client.attach(create_result['Id'], stream=True, logs=True):
        sys.stdout.write(line)

    '''
    TODO: if ^C will be pressed before this, container will not be stopped and removed.
    We need to handle signal here.
    '''
    docker_client.stop(create_result['Id'])
    docker_client.remove_container(create_result['Id'])

def run_docker_container(docker_client, docker_hash, service, args):
    """
    Run Docker container by image hash with args set in command line.
    Function prints the output of run command in stdout.
    """
    marathon_config_raw = read_marathon_config()

    volumes = list()

    for volume in marathon_config_raw['docker_volumes']:
        volumes.append('%s:%s:%s' % (volume['hostPath'], volume['containerPath'], volume['mode'].lower()))

    marathon_config = dict()

    marathon_config['cluster'] = marathon_config_raw['cluster']
    marathon_config['volumes'] = volumes

    service_manifest = read_service_config(service, args.instance, marathon_config['cluster'])

    if args.cmd:
        command = shlex.split(args.cmd)
    else:
        command = get_args(service_manifest)

    if args.interactive:
        run_docker_container_interactive(
            service,
            args.instance,
            docker_hash,
            volumes,
            command,
            service_manifest
        )
    else:
        run_docker_container_non_interactive(
            docker_client,
            service,
            args.instance,
            docker_hash,
            volumes,
            command,
            service_manifest
        )

def build_docker_container(docker_client, args):
    """
    Build Docker container from Dockerfile in the current directory or
    specified in command line args. Resulting image hash.
    """
    result = ''

    dockerfile_path = os.getcwd()

    sys.stdout.write('Building container from Dockerfile in %s\n' % dockerfile_path)

    for line in docker_client.build(path=dockerfile_path, tag='latest'):
        line_dict = json.loads(line)

        stream_line = line_dict.get('stream')

        if args.verbose:
            sys.stdout.write(stream_line)

        if stream_line.startswith('Successfully built '):
            """Strip the beginning of a string and \n in the end."""
            result = stream_line[len('Successfully built '):]
            result = result[:len(result) - 1]

    return result

def paasta_test_run(args):
    service = figure_out_service_name(args)

    base_docker_url = os.environ.get('DOCKER_HOST', 'unix://var/run/docker.sock')

    docker_client = Client(base_url=base_docker_url)

    try:
        docker_hash = build_docker_container(docker_client, args)
    except errors.APIError as e:
        sys.stderr.write('Can\'t build Docker container. Error: %s\n' % str(e))
        sys.exit(1)

    try:
        run_docker_container(docker_client, docker_hash, service, args)
    except errors.APIError as e:
        sys.stderr.write('Can\'t run Docker container. Error: %s\n' % str(e))
        sys.exit(1)
<|endoftext|>


<commit_before>"""
Usage:

    ./manage.py import_operators < NOC_db.csv
"""

from busstops.management.import_from_csv import ImportFromCSVCommand
from busstops.models import Operator, Region

class Command(ImportFromCSVCommand):

    @staticmethod
    def get_region(region_id):
        if region_id in ('ADMIN', 'Admin', ''):
            region_id = 'GB'
        elif region_id in ('SC', 'YO', 'WA', 'LO'):
            region_id = region_id[0]

        return Region.objects.get(id=region_id)

    @classmethod
    def handle_row(cls, row):
        "Given a CSV row (a list), returns an Operator object"
        operator_id = row['NOCCODE'].replace('=', '')

        if operator_id == 'TVSR':
            return None

        region = cls.get_region(row['TLRegOwn'])

        if row['OperatorPublicName'] in ('First', 'Arriva', 'Stagecoach') \
            or row['OperatorPublicName'].startswith('inc.') \
            or row['OperatorPublicName'].startswith('formerly'):
            name = row['RefNm']
        else:
            name = row['OperatorPublicName']

        name = name.replace('\'', u'\u2019') # Fancy apostrophe

        operator = Operator.objects.update_or_create(
            id=operator_id,
            defaults=dict(
                name=name.strip(),
                vehicle_mode=row['Mode'].lower().replace('ct operator', 'community transport').replace('drt', 'demand responsive transport'),
                region=region,
            )
        )
        return operator
<commit_msg>Fix some operators with blank names<commit_after>"""
Usage:

    ./manage.py import_operators < NOC_db.csv
"""

from busstops.management.import_from_csv import ImportFromCSVCommand
from busstops.models import Operator, Region

class Command(ImportFromCSVCommand):

    @staticmethod
    def get_region(region_id):
        if region_id in ('ADMIN', 'Admin', ''):
            region_id = 'GB'
        elif region_id in ('SC', 'YO', 'WA', 'LO'):
            region_id = region_id[0]

        return Region.objects.get(id=region_id)

    @classmethod
    def handle_row(cls, row):
        "Given a CSV row (a list), returns an Operator object"
        operator_id = row['NOCCODE'].replace('=', '')

        if operator_id == 'TVSR':
            return None

        region = cls.get_region(row['TLRegOwn'])

        if row['OperatorPublicName'] in ('First', 'Arriva', 'Stagecoach') \
            or row['OperatorPublicName'].startswith('inc.') \
            or row['OperatorPublicName'].startswith('formerly'):
            name = row['RefNm']
        elif row['OperatorPublicName'] != '':
            name = row['OperatorPublicName']
        else:
            name = row['OpNm']

        name = name.replace('\'', u'\u2019') # Fancy apostrophe

        operator = Operator.objects.update_or_create(
            id=operator_id,
            defaults=dict(
                name=name.strip(),
                vehicle_mode=row['Mode'].lower().replace('ct operator', 'community transport').replace('drt', 'demand responsive transport'),
                region=region,
            )
        )
        return operator
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
#############################################
# Copyright (C) 2018 Z        (User:ZxxZxxZ)
#                    Reza1615 (User:Reza1615)
#                    Huji     (User:Huji)
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#############################################

from __future__ import absolute_import, unicode_literals

import re
import time
import webcitation
import pywikibot

from pywikibot import pagegenerators
from pywikibot.bot import (
    SingleSiteBot, ExistingPageBot, NoRedirectPageBot, AutomaticTWSummaryBot)
from datetime import (
    timedelta, datetime)

class WebCiteBot(
    SingleSiteBot,
    ExistingPageBot,
    NoRedirectPageBot
):

    deadlink = []

    def __init__(self, generator, **kwargs):
        # Edit summary
        self.availableOptions.update({
            'summary': u'ربات: [[وپ:پایدار|بایگانی منابع برخط]]',
        })

        # Call constructor of the super class
        super(WebCiteBot, self).__init__(site=True, **kwargs)

       # Assign the generator to the bot
        self.generator = generator

        # Save edits without asking
        self.options['always'] = True

        # Global counter for all archived links
        self.counter = 0

        # Maximum number of links that can be archived each day
        self.availableOptions.update({
            'maxlinks': 99,
        })

    def persianDigits(self, s):
        return str(s).replace(u'0',u'۰') \
            .replace(u'1',u'۱') \
            .replace(u'2',u'۲') \
            .replace(u'3',u'۳') \
            .replace(u'4',u'۴') \
            .replace(u'5',u'۵') \
            .replace(u'6',u'۶') \
            .replace(u'7',u'۷') \
            .replace(u'8',u'۸') \
            .replace(u'9',u'۹')

    def persianDate(self):
        monthNames = [
            u"ژانویه",
            u"فوریه",
            u"مارس",
            u"آوریل",
            u"مه",
            u"ژوئن",
            u"ژوئیه",
            u"اوت",
            u"سپتامبر",
            u"اکتبر",
            u"نوامبر",
            u"دسامبر"
            ]

        today = datetime.today()
        y = today.year
        m = monthNames[today.month - 1]
        d = today.day

        d = self.persianDigits(d)
        y = self.persianDigits(y)

        return '%s %s %s' % (d, m, y)

    def archive(self, URL):
        pywikibot.output("Requesting URL %s" % URL)
        try:
            archiveURL = webcitation.capture(URL)
        except:
            pywikibot.output(u"\03{red}Archive attempt failed!\03{default}")
            return False
        pywikibot.output(u"\03{lightgreen}Successfully archived at %s\03{default}" % archiveURL)
        self.counter += 1
        return archiveURL

    def treat_page(self):
        if (self.current_page.namespace() != 0):
            pywikibot.output("Not an article; skipped...")
            return True

        text = self.current_page.text

        # Persian Citations
        persianDate = self.persianDate()
        faCitationPattern = ur'\{\{\s*یادکرد(?:\{\{.*?\}\}|.)*?\}\}'
        faArchivePattern = ur'\| *(پیوند بایگانی|نشانی بایگانی) *= *[^ |}]'
        faBlankArchivePattern = ur'\| *(پیوند بایگانی|نشانی بایگانی|تاریخ بایگانی) *= *(?=[|}])'
        faUrlPattern = ur'\| *(نشانی|پیوند) *= *([^|]+) *\|'
        faCitations = re.findall(faCitationPattern, text, re.S)
        if not faCitations:
            pywikibot.output(u"\03{lightpurple}No Persian citations!\03{default}")
        else:
            for citation in faCitations:
                if not re.findall(faArchivePattern, citation, re.S):
                    # Not archived already
                    url = re.findall(faUrlPattern, citation, re.S)
                    if url and url[0][1].strip() != "":
                        # Try archiving it
                        if self.counter + 1 > self.getOption('maxlinks'):
                            pywikibot.output(u"\03{red}WebCitation quota is used up!\03{default}")
                            continue

                        arc = self.archive(url[0][1])
                        if not arc:
                            # Archiving failed
                            continue
                        # Remove any blank archiveurl parameters
                        newCitation = re.sub(faBlankArchivePattern, "", citation)
                        # Add archiveurl and archivedate
                        newParams = "| پیوند بایگانی = %s | تاریخ بایگانی = %s }}" % (arc, persianDate)
                        newCitation = re.sub("}}", newParams, newCitation)
                        text = text.replace(citation, newCitation)

        # English Citations
        englishDate = datetime.today().strftime("%-d %B %Y")
        enCitationPattern = ur'\{\{\s*cite(?:\{\{.*?\}\}|.)*?\}\}'
        enArchivePattern = ur'\| *(archiveurl|archive-url) *= *[^ |}]'
        enBlankArchivePattern = ur'\| *(archiveurl|archive-url|archivedate|archive-date) *= *(?=[|}])'
        enUrlPattern = ur'\| *(url) *= *([^|]+) *\|'
        enCitations = re.findall(enCitationPattern, text, re.S)
        if not enCitations:
            pywikibot.output(u"\03{lightpurple}No English citations!\03{default}")
        else:
            for citation in enCitations:
                if not re.findall(enArchivePattern, citation, re.S):
                    # Not archived already
                    url = re.findall(enUrlPattern, citation, re.S)
                    if url and url[0][1].strip() != "":
                        # Try archiving it
                        if self.counter + 1 > self.getOption('maxlinks'):
                            pywikibot.output(u"\03{red}WebCitation quota is used up!\03{default}")
                            continue

                        arc = self.archive(url[0][1])
                        if not arc:
                            # Archiving failed
                            continue
                        # Remove any blank archiveurl parameters
                        newCitation = re.sub(enBlankArchivePattern, "", citation)
                        # Add archiveurl and archivedate
                        newParams = "| archive-url = %s | archive-date = %s }}" % (arc, englishDate)
                        newCitation = re.sub("}}", newParams, newCitation)
                        text = text.replace(citation, newCitation)

        self.put_current(text, summary=self.getOption('summary'))

def main(*args):
    options = {}
    local_args = pywikibot.handle_args(args)
    genFactory = pagegenerators.GeneratorFactory()

    for arg in local_args:
        # Catch the pagegenerators options
        if genFactory.handleArg(arg):
            continue  # nothing to do here
        # Take the remaining options as booleans.
        else:
            pywikibot.output('Unknown argument(s)!')

    gen = genFactory.getCombinedGenerator()
    if gen:
        gen = pagegenerators.PreloadingGenerator(gen)
        bot = WebCiteBot(gen, **options)
        bot.run()
        return True
    else:
        pywikibot.bot.suggest_help(missing_generator=True)
        return False

if __name__ == "__main__":
    main()

<commit_msg>For some reason, ur does not work on Labs, but u does<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
#############################################
# Copyright (C) 2018 Z        (User:ZxxZxxZ)
#                    Reza1615 (User:Reza1615)
#                    Huji     (User:Huji)
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#############################################

from __future__ import absolute_import, unicode_literals

import re
import time
import webcitation
import pywikibot

from pywikibot import pagegenerators
from pywikibot.bot import (
    SingleSiteBot, ExistingPageBot, NoRedirectPageBot, AutomaticTWSummaryBot)
from datetime import (
    timedelta, datetime)

class WebCiteBot(
    SingleSiteBot,
    ExistingPageBot,
    NoRedirectPageBot
):

    deadlink = []

    def __init__(self, generator, **kwargs):
        # Edit summary
        self.availableOptions.update({
            'summary': u'ربات: [[وپ:پایدار|بایگانی منابع برخط]]',
        })

        # Call constructor of the super class
        super(WebCiteBot, self).__init__(site=True, **kwargs)

       # Assign the generator to the bot
        self.generator = generator

        # Save edits without asking
        self.options['always'] = True

        # Global counter for all archived links
        self.counter = 0

        # Maximum number of links that can be archived each day
        self.availableOptions.update({
            'maxlinks': 99,
        })

    def persianDigits(self, s):
        return str(s).replace(u'0',u'۰') \
            .replace(u'1',u'۱') \
            .replace(u'2',u'۲') \
            .replace(u'3',u'۳') \
            .replace(u'4',u'۴') \
            .replace(u'5',u'۵') \
            .replace(u'6',u'۶') \
            .replace(u'7',u'۷') \
            .replace(u'8',u'۸') \
            .replace(u'9',u'۹')

    def persianDate(self):
        monthNames = [
            u"ژانویه",
            u"فوریه",
            u"مارس",
            u"آوریل",
            u"مه",
            u"ژوئن",
            u"ژوئیه",
            u"اوت",
            u"سپتامبر",
            u"اکتبر",
            u"نوامبر",
            u"دسامبر"
            ]

        today = datetime.today()
        y = today.year
        m = monthNames[today.month - 1]
        d = today.day

        d = self.persianDigits(d)
        y = self.persianDigits(y)

        return '%s %s %s' % (d, m, y)

    def archive(self, URL):
        pywikibot.output("Requesting URL %s" % URL)
        try:
            archiveURL = webcitation.capture(URL)
        except:
            pywikibot.output(u"\03{red}Archive attempt failed!\03{default}")
            return False
        pywikibot.output(u"\03{lightgreen}Successfully archived at %s\03{default}" % archiveURL)
        self.counter += 1
        return archiveURL

    def treat_page(self):
        if (self.current_page.namespace() != 0):
            pywikibot.output("Not an article; skipped...")
            return True

        text = self.current_page.text

        # Persian Citations
        persianDate = self.persianDate()
        faCitationPattern = u'\{\{\s*یادکرد(?:\{\{.*?\}\}|.)*?\}\}'
        faArchivePattern = u'\| *(پیوند بایگانی|نشانی بایگانی) *= *[^ |}]'
        faBlankArchivePattern = u'\| *(پیوند بایگانی|نشانی بایگانی|تاریخ بایگانی) *= *(?=[|}])'
        faUrlPattern = u'\| *(نشانی|پیوند) *= *([^|]+) *\|'
        faCitations = re.findall(faCitationPattern, text, re.S)
        if not faCitations:
            pywikibot.output(u"\03{lightpurple}No Persian citations!\03{default}")
        else:
            for citation in faCitations:
                if not re.findall(faArchivePattern, citation, re.S):
                    # Not archived already
                    url = re.findall(faUrlPattern, citation, re.S)
                    if url and url[0][1].strip() != "":
                        # Try archiving it
                        if self.counter + 1 > self.getOption('maxlinks'):
                            pywikibot.output(u"\03{red}WebCitation quota is used up!\03{default}")
                            continue

                        arc = self.archive(url[0][1])
                        if not arc:
                            # Archiving failed
                            continue
                        # Remove any blank archiveurl parameters
                        newCitation = re.sub(faBlankArchivePattern, "", citation)
                        # Add archiveurl and archivedate
                        newParams = "| پیوند بایگانی = %s | تاریخ بایگانی = %s }}" % (arc, persianDate)
                        newCitation = re.sub("}}", newParams, newCitation)
                        text = text.replace(citation, newCitation)

        # English Citations
        englishDate = datetime.today().strftime("%-d %B %Y")
        enCitationPattern = u'\{\{\s*cite(?:\{\{.*?\}\}|.)*?\}\}'
        enArchivePattern = u'\| *(archiveurl|archive-url) *= *[^ |}]'
        enBlankArchivePattern = u'\| *(archiveurl|archive-url|archivedate|archive-date) *= *(?=[|}])'
        enUrlPattern = u'\| *(url) *= *([^|]+) *\|'
        enCitations = re.findall(enCitationPattern, text, re.S)
        if not enCitations:
            pywikibot.output(u"\03{lightpurple}No English citations!\03{default}")
        else:
            for citation in enCitations:
                if not re.findall(enArchivePattern, citation, re.S):
                    # Not archived already
                    url = re.findall(enUrlPattern, citation, re.S)
                    if url and url[0][1].strip() != "":
                        # Try archiving it
                        if self.counter + 1 > self.getOption('maxlinks'):
                            pywikibot.output(u"\03{red}WebCitation quota is used up!\03{default}")
                            continue

                        arc = self.archive(url[0][1])
                        if not arc:
                            # Archiving failed
                            continue
                        # Remove any blank archiveurl parameters
                        newCitation = re.sub(enBlankArchivePattern, "", citation)
                        # Add archiveurl and archivedate
                        newParams = "| archive-url = %s | archive-date = %s }}" % (arc, englishDate)
                        newCitation = re.sub("}}", newParams, newCitation)
                        text = text.replace(citation, newCitation)

        self.put_current(text, summary=self.getOption('summary'))

def main(*args):
    options = {}
    local_args = pywikibot.handle_args(args)
    genFactory = pagegenerators.GeneratorFactory()

    for arg in local_args:
        # Catch the pagegenerators options
        if genFactory.handleArg(arg):
            continue  # nothing to do here
        # Take the remaining options as booleans.
        else:
            pywikibot.output('Unknown argument(s)!')

    gen = genFactory.getCombinedGenerator()
    if gen:
        gen = pagegenerators.PreloadingGenerator(gen)
        bot = WebCiteBot(gen, **options)
        bot.run()
        return True
    else:
        pywikibot.bot.suggest_help(missing_generator=True)
        return False

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>"""Leetcode 160. Intersection of Two Linked Lists
Easy

URL: https://leetcode.com/problems/intersection-of-two-linked-lists/

Write a program to find the node at which the intersection of
two singly linked lists begins.

For example, the following two linked lists:

begin to intersect at node c1.

Example 1:
Input: intersectVal = 8, listA = [4,1,8,4,5], listB = [5,0,1,8,4,5],
skipA = 2, skipB = 3
Output: Reference of the node with value = 8
Input Explanation: The intersected node's value is 8
(note that this must not be 0 if the two lists intersect).
From the head of A, it reads as [4,1,8,4,5]. From the head of B,
it reads as [5,0,1,8,4,5].
There are 2 nodes before the intersected node in A;
There are 3 nodes before the intersected node in B.

Example 2:
Input: intersectVal = 2, listA = [0,9,1,2,4], listB = [3,2,4], skipA = 3, skipB = 1
Output: Reference of the node with value = 2
Input Explanation: The intersected node's value is 2
(note that this must not be 0 if the two lists intersect).
From the head of A, it reads as [0,9,1,2,4].
From the head of B, it reads as [3,2,4].
There are 3 nodes before the intersected node in A;
There are 1 node before the intersected node in B.

Example 3:
Input: intersectVal = 0, listA = [2,6,4], listB = [1,5], skipA = 3, skipB = 2
Output: null
Input Explanation: From the head of A, it reads as [2,6,4].
From the head of B, it reads as [1,5].
Since the two lists do not intersect, intersectVal must be 0,
while skipA and skipB can be arbitrary values.
Explanation: The two lists do not intersect, so return null.

Notes:
- If the two linked lists have no intersection at all, return null.
- The linked lists must retain their original structure after the function returns.
- You may assume there are no cycles anywhere in the entire linked structure.
- Your code should preferably run in O(n) time and use only O(1) memory.
"""

# Definition for singly-linked list.
class ListNode(object):
    def __init__(self, val):
        self.val = val
        self.next = None

class Solution(object):
    def getIntersectionNode(self, headA, headB):
        """
        :type head1, head1: ListNode
        :rtype: ListNode
        """
        pass

def main():
    pass

if __name__ == '__main__':
    main()
<commit_msg>Complete sol with two pointers<commit_after>"""Leetcode 160. Intersection of Two Linked Lists
Easy

URL: https://leetcode.com/problems/intersection-of-two-linked-lists/

Write a program to find the node at which the intersection of
two singly linked lists begins.

For example, the following two linked lists:

begin to intersect at node c1.

Example 1:
Input: intersectVal = 8, listA = [4,1,8,4,5], listB = [5,0,1,8,4,5],
skipA = 2, skipB = 3
Output: Reference of the node with value = 8
Input Explanation: The intersected node's value is 8
(note that this must not be 0 if the two lists intersect).
From the head of A, it reads as [4,1,8,4,5]. From the head of B,
it reads as [5,0,1,8,4,5].
There are 2 nodes before the intersected node in A;
There are 3 nodes before the intersected node in B.

Example 2:
Input: intersectVal = 2, listA = [0,9,1,2,4], listB = [3,2,4], skipA = 3, skipB = 1
Output: Reference of the node with value = 2
Input Explanation: The intersected node's value is 2
(note that this must not be 0 if the two lists intersect).
From the head of A, it reads as [0,9,1,2,4].
From the head of B, it reads as [3,2,4].
There are 3 nodes before the intersected node in A;
There are 1 node before the intersected node in B.

Example 3:
Input: intersectVal = 0, listA = [2,6,4], listB = [1,5], skipA = 3, skipB = 2
Output: null
Input Explanation: From the head of A, it reads as [2,6,4].
From the head of B, it reads as [1,5].
Since the two lists do not intersect, intersectVal must be 0,
while skipA and skipB can be arbitrary values.
Explanation: The two lists do not intersect, so return null.

Notes:
- If the two linked lists have no intersection at all, return null.
- The linked lists must retain their original structure after the function returns.
- You may assume there are no cycles anywhere in the entire linked structure.
- Your code should preferably run in O(n) time and use only O(1) memory.
"""

# Definition for singly-linked list.
class ListNode(object):
    def __init__(self, val):
        self.val = val
        self.next = None

class SolutionTwoPointers(object):
    def getIntersectionNode(self, headA, headB):
        """
        :type head1, head1: ListNode
        :rtype: ListNode
        """
        if not headA or not headB:
            return None

        # Two pointer method: Check current nodes are the same or not.
        currentA = headA
        currentB = headB

        while currentA != currentB:
            if currentA:
                # If not, visit next nodes.
                currentA = currentA.next
            else:
                # When arrive at the end, visit the head of another list.
                currentA = headB

            if currentB:
                currentB = currentB.next
            else:
                currentB = headA

        return currentA

def main():
    # Input: intersectVal = 8, listA = [4,1,8,4,5], listB = [5,0,1,8,4,5],
    # Output: Reference of the node with value = 8
    headA = ListNode(4)
    headA.next = ListNode(1)
    headB = ListNode(5)
    headB.next = ListNode(0)
    headB.next.next = ListNode(1)
    intersectNode1 = ListNode(8)
    intersectNode2 = ListNode(4)
    intersectNode3 = ListNode(5)

    headA.next.next = intersectNode1
    headA.next.next.next = intersectNode2
    headA.next.next.next.next = intersectNode3
    headB.next.next = intersectNode1
    headB.next.next.next = intersectNode2
    headB.next.next.next.next = intersectNode3

    print SolutionTwoPointers().getIntersectionNode(headA, headB).val

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from .common import InfoExtractor
from ..compat import compat_str
from ..utils import (
    int_or_none,
    unified_timestamp,
)

class BeegIE(InfoExtractor):
    _VALID_URL = r'https?://(?:www\.)?beeg\.(?:com|porn(?:/video)?)/(?P<id>\d+)'
    _TESTS = [{
        'url': 'http://beeg.com/5416503',
        'md5': 'a1a1b1a8bc70a89e49ccfd113aed0820',
        'info_dict': {
            'id': '5416503',
            'ext': 'mp4',
            'title': 'Sultry Striptease',
            'description': 'md5:d22219c09da287c14bed3d6c37ce4bc2',
            'timestamp': 1391813355,
            'upload_date': '20140207',
            'duration': 383,
            'tags': list,
            'age_limit': 18,
        }
    }, {
        'url': 'https://beeg.porn/video/5416503',
        'only_matching': True,
    }, {
        'url': 'https://beeg.porn/5416503',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        video_id = self._match_id(url)

        webpage = self._download_webpage(url, video_id)

        beeg_version = self._search_regex(
            r'beeg_version\s*=\s*([\da-zA-Z_-]+)', webpage, 'beeg version',
            default='1546225636701')

        for api_path in ('', 'api.'):
            video = self._download_json(
                'https://%sbeeg.com/api/v6/%s/video/%s'
                % (api_path, beeg_version, video_id), video_id,
                fatal=api_path == 'api.')
            if video:
                break

        formats = []
        for format_id, video_url in video.items():
            if not video_url:
                continue
            height = self._search_regex(
                r'^(\d+)[pP]$', format_id, 'height', default=None)
            if not height:
                continue
            formats.append({
                'url': self._proto_relative_url(
                    video_url.replace('{DATA_MARKERS}', 'data=pc_XX__%s_0' % beeg_version), 'https:'),
                'format_id': format_id,
                'height': int(height),
            })
        self._sort_formats(formats)

        title = video['title']
        video_id = compat_str(video.get('id') or video_id)
        display_id = video.get('code')
        description = video.get('desc')
        series = video.get('ps_name')

        timestamp = unified_timestamp(video.get('date'))
        duration = int_or_none(video.get('duration'))

        tags = [tag.strip() for tag in video['tags'].split(',')] if video.get('tags') else None

        return {
            'id': video_id,
            'display_id': display_id,
            'title': title,
            'description': description,
            'series': series,
            'timestamp': timestamp,
            'duration': duration,
            'tags': tags,
            'formats': formats,
            'age_limit': self._rta_search(webpage),
        }
<commit_msg>[beeg] Add support for api/v6 v2 URLs (closes #21511)<commit_after>from __future__ import unicode_literals

from .common import InfoExtractor
from ..compat import (
    compat_str,
    compat_urlparse,
)
from ..utils import (
    int_or_none,
    unified_timestamp,
)

class BeegIE(InfoExtractor):
    _VALID_URL = r'https?://(?:www\.)?beeg\.(?:com|porn(?:/video)?)/(?P<id>\d+)'
    _TESTS = [{
        # api/v6 v1
        'url': 'http://beeg.com/5416503',
        'md5': 'a1a1b1a8bc70a89e49ccfd113aed0820',
        'info_dict': {
            'id': '5416503',
            'ext': 'mp4',
            'title': 'Sultry Striptease',
            'description': 'md5:d22219c09da287c14bed3d6c37ce4bc2',
            'timestamp': 1391813355,
            'upload_date': '20140207',
            'duration': 383,
            'tags': list,
            'age_limit': 18,
        }
    }, {
        # api/v6 v2
        'url': 'https://beeg.com/1941093077?t=911-1391',
        'only_matching': True,
    }, {
        'url': 'https://beeg.porn/video/5416503',
        'only_matching': True,
    }, {
        'url': 'https://beeg.porn/5416503',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        video_id = self._match_id(url)

        webpage = self._download_webpage(url, video_id)

        beeg_version = self._search_regex(
            r'beeg_version\s*=\s*([\da-zA-Z_-]+)', webpage, 'beeg version',
            default='1546225636701')

        qs = compat_urlparse.parse_qs(compat_urlparse.urlparse(url).query)
        t = qs.get('t', [''])[0].split('-')
        if len(t) > 1:
            query = {
                'v': 2,
                's': t[0],
                'e': t[1],
            }
        else:
            query = {'v': 1}

        for api_path in ('', 'api.'):
            video = self._download_json(
                'https://%sbeeg.com/api/v6/%s/video/%s'
                % (api_path, beeg_version, video_id), video_id,
                fatal=api_path == 'api.', query=query)
            if video:
                break

        formats = []
        for format_id, video_url in video.items():
            if not video_url:
                continue
            height = self._search_regex(
                r'^(\d+)[pP]$', format_id, 'height', default=None)
            if not height:
                continue
            formats.append({
                'url': self._proto_relative_url(
                    video_url.replace('{DATA_MARKERS}', 'data=pc_XX__%s_0' % beeg_version), 'https:'),
                'format_id': format_id,
                'height': int(height),
            })
        self._sort_formats(formats)

        title = video['title']
        video_id = compat_str(video.get('id') or video_id)
        display_id = video.get('code')
        description = video.get('desc')
        series = video.get('ps_name')

        timestamp = unified_timestamp(video.get('date'))
        duration = int_or_none(video.get('duration'))

        tags = [tag.strip() for tag in video['tags'].split(',')] if video.get('tags') else None

        return {
            'id': video_id,
            'display_id': display_id,
            'title': title,
            'description': description,
            'series': series,
            'timestamp': timestamp,
            'duration': duration,
            'tags': tags,
            'formats': formats,
            'age_limit': self._rta_search(webpage),
        }
<|endoftext|>