#!/usr/bin/python

# Copyright (C) 2005-2017 Canonical Ltd.
# Authors: Kees Cook <kees@ubuntu.com>
#          Jamie Strandboge <jamie@ubuntu.com>
#          Marc Deslauriers <marc.deslauriers@canonical.com>
#
# This script is distributed under the terms and conditions of the GNU General
# Public License, Version 3 or later. See http://www.gnu.org/copyleft/gpl.html
# for details.

from __future__ import print_function

import apt_pkg
import os
import re
import subprocess
import sys
from ivulncheck_api.lib.uct import cve_lib

from ivulncheck_api.lib.config_api import Configuration as conf

try:
    from configobj import ConfigObj
except:
    # Dapper lacks this class, so reimplement it quickly
    class ConfigObj(dict):
        def __init__(self, filepath):
            with open(filepath) as f:
                for line in f.readlines():
                    line = line.strip()
                    if line.startswith('#') or len(line)==0:
                        continue
                    name, stuff = line.strip().split('=',1)
                    self[name] = eval(stuff)

        def __attr__(self, name):
            return self.stuff[name]

apt_pkg.init_system();

def version_compare(one,two):
    return apt_pkg.version_compare(one, two)

def _find_sources(pockets=None,releases=None):
    config = ConfigObj(os.path.expanduser(conf.getUCTConfFile()))
    if 'packages_mirror' in config and 'partner_mirror' in config:
        cve_lib.check_mirror_timestamp(config)
        return _find_from_mirror(config['packages_mirror'],
                                 config['partner_mirror'],
                                 'sources',
                                 pockets=pockets,
                                 releases=releases)
    else:
        return _find_sources_from_apt(pockets=pockets,releases=releases)

def _find_packages(pockets=None,releases=None):
    config = ConfigObj(os.path.expanduser(conf.getUCTConfFile()))
    if 'packages_mirror' in config and 'partner_mirror' in config:
        cve_lib.check_mirror_timestamp(config)
        return _find_from_mirror(config['packages_mirror'],
                                 config['partner_mirror'],
                                 'packages',
                                 pockets=pockets,
                                 releases=releases)
    else:
        raise Exception("TODO: implement _find_packages_from_apt()")
        # return _find_packages_from_apt()

def load_debian(basedir, data_type='sources'):
    if data_type not in ['sources', 'packages']:
        raise ValueError("'data_type' should be either 'sources' or 'packages'")

    debian_sources = dict()
    for section in ['main','contrib','non-free']:
        rel = 'testing'
        if data_type == 'sources':
            load_sources_collection( (os.path.join(basedir,'dists', rel, section, 'source', 'Sources.gz'), rel, '', section), debian_sources )
    return debian_sources

def _find_from_mirror(ubuntu, canonical, data_type, arch='amd64', pockets=None, releases=None):
    if data_type not in ['sources', 'packages']:
        raise ValueError("'data_type' should be either 'sources' or 'packages'")

    collection = []
    errors = False
    missing = ""

    if pockets is None:
        pockets = ['','-updates','-security']
    if releases is None:
        releases = cve_lib.releases

    # Dev release may not be opened
    if releases == ['']:
        return collection

    for rel in releases:
        if rel in cve_lib.eol_releases:
            continue
        if '/' in rel:  # we handle ppa overlays differently
            continue
        # free
        for pocket in pockets:
            for section in ['main','restricted','universe','multiverse']:
                if data_type == 'sources':
                    fn = os.path.join(ubuntu, 'dists', rel + pocket, section,
                                      'source', 'Sources')
                else:
                    fn = os.path.join(ubuntu, 'dists', rel + pocket, section,
                                      'binary-%s' % arch, 'Packages')

                if not os.path.exists(fn):
                    fn += '.gz'
                    if not os.path.exists(fn):
                        missing += " %s\n" % (fn)
                        errors = True
                        continue
                name = pocket
                if name.startswith('-'):
        	        name = name[1:]
                else:
                    name = ''
                collection += [(fn, rel, name, section)]
        # partner
        pocket = ''
        section = 'partner'

        if data_type == 'sources':
            fn = os.path.join(canonical, 'dists', rel + pocket, section,
                              'source', 'Sources')
        else:
            fn = os.path.join(canonical, 'dists', rel + pocket, section,
                              'binary-%s' % arch, 'Packages')

        if not os.path.exists(fn):
            fn += '.gz'
            if not os.path.exists(fn):
                # Only warn about missing partner for devel release
                if rel == cve_lib.devel_release:
                    prefreeze = os.path.join(canonical, 'dists', '%s-series' % rel[0])
                    if not os.path.isdir(prefreeze):
                        print(sys.stderr, "WARNING: missing partner mirror element: %s" % (fn), file=sys.stderr)
                    #else:
                    #    print("WARNING: found prefreeze element: %s" % (prefreeze), file=sys.stderr)
                else:
                    missing += " %s\n" % (fn)
                    errors = True
                continue
        collection += [(fn, rel, '', section)]

    if errors:
        raise NameError("Missing mirror elements:\n" + missing)

    return collection


def _find_sources_from_apt(pockets=None,releases=None):
  collection = []

  if pockets is None:
      pockets = ['','-updates','-security']
  if releases is None:
      releases = cve_lib.releases

  saw = dict()
  lists = '/var/lib/apt/lists'
  for f in os.listdir(lists):
    if not f.endswith('_source_Sources') and not '-commercial_main_binary-' in f:
        continue
    parts = f.split('_')
    parts.pop() # _Sources
    parts.pop() # _source
    section = parts.pop() # _main
    release_real = parts.pop() # _dapper
    saw.setdefault(release_real,True)
    tmp = release_real.split('-')
    release = tmp[0]
    if len(tmp) > 1:
    	pocket = tmp[1]
    else:
        pocket = ''
    collection += [(os.path.join(lists,f), release, pocket, section)]

  # Validate all the sources are available
  errors = False
  missing = ""
  for rel in releases:
    if rel in cve_lib.eol_releases:
        continue
    if '/' in rel:  # we handle ppa overlays differently
        continue
    for pocket in pockets:
        if rel+pocket not in saw:
            missing += " deb-src http://archive.ubuntu.com/ubuntu %s%s main restricted universe multiverse\n" % (rel,pocket)
            errors = True
    for pocket in ['-commercial']:
        if rel < 'gutsy' and rel+pocket not in saw:
            missing += " deb http://archive.canonical.com/ubuntu %s%s main\n" % (rel,pocket)
            errors = True
  if errors:
    raise NameError("Missing /etc/apt/sources.list lines:\n%s" % (missing))

  return collection

# release -> pkg -> dict( 'section', 'pocket', 'version' )
def load(data_type='sources',pockets=None,releases=None):
    if data_type not in ['sources', 'packages']:
        raise ValueError("'data_type' should be either 'sources' or 'packages'")

    map = dict()
    if data_type == 'sources':
        for item in _find_sources(pockets=pockets,releases=releases):
            load_sources_collection(item, map)
    else:
        for item in _find_packages(pockets=pockets,releases=releases):
            load_packages_collection(item, map)

    ppa_overlay = load_ppa()
    for item in ppa_overlay:
        map[item] = ppa_overlay[item]
    return map

def _get_apt_tags(tagfile):
    tags = None
    if tagfile.endswith('.gz'):
        tags = subprocess.Popen(['/bin/gunzip','-c',tagfile], stdout=subprocess.PIPE).stdout
    elif tagfile.endswith('.bz2'):
        tags = subprocess.Popen(['/bin/bunzip2','-c',tagfile], stdout=subprocess.PIPE).stdout
    else:
        tags = open(tagfile)

    return tags

def load_sources_collection(item, map):
    tagfile, release, pocket, section = item

    parser = apt_pkg.TagFile(_get_apt_tags(tagfile))
    while parser.step():
        pkg = parser.section['Package']
        map.setdefault(release,dict()).setdefault(pkg, {'section': 'unset', 'version': '~', 'pocket': 'unset' })
        map[release][pkg]['section'] = section
        if not pocket:
            map[release][pkg]['release_version'] = parser.section['Version']
        if apt_pkg.version_compare(parser.section['Version'],map[release][pkg]['version'])>0:
            map[release][pkg]['pocket'] = pocket
            map[release][pkg]['version'] = parser.section['Version']
            map[release][pkg]['binaries'] = parser.section['Binary'].split(', ')

    return map

def load_packages_collection(item, map):
    tagfile, release, pocket, section = item

    parser = apt_pkg.TagFile(_get_apt_tags(tagfile))
    while parser.step():
        pkg = parser.section['Package']
        map.setdefault(release,dict()).setdefault(pkg, {'section': 'unset', 'version': '~', 'pocket': 'unset' })
        map[release][pkg]['section'] = section
        if not pocket:
            map[release][pkg]['release_version'] = parser.section['Version']
        if apt_pkg.version_compare(parser.section['Version'],map[release][pkg]['version'])>0:
            map[release][pkg]['pocket'] = pocket
            map[release][pkg]['version'] = parser.section['Version']

            if 'Source' in parser.section:
                map[release][pkg]['source'] = parser.section['Source'].split()[0]
            else:
                map[release][pkg]['source'] = parser.section['Package']

            if 'Built-Using' in parser.section:
                map[release][pkg]['built-using'] = parser.section['Built-Using'].split(', ')

    return map

def load_built_using_collection(pmap, releases=None, component=None):
    built_using = dict()

    for rel in pmap.keys():
        if releases is not None and rel not in releases:
            continue
        if '/' in rel:  # we handle ppa overlays differently
            continue

        for pkg in pmap[rel]:
            if 'built-using' in pmap[rel][pkg]:
                # Built-Using for a binary in the Packages file lists the
                # originating source package of the embedded binary
                section = pmap[rel][pkg]['section']
                if component is not None and section != component:
                    continue

                pocket = rel
                if pmap[rel][pkg]['pocket'] != '':
                    pocket += "-%s" % pmap[rel][pkg]['pocket']

                for (s, c, v) in map(lambda x: x.split(' ', 3),
                                     pmap[rel][pkg]['built-using']):
                    v = v.rstrip(')')
                    if s not in built_using:
                        built_using[s] = dict()
                    if v not in built_using[s]:
                        built_using[s][v] = dict()
                    if section not in built_using[s][v]:
                        built_using[s][v][section] = dict()
                    if pocket not in built_using[s][v][section]:
                        built_using[s][v][section][pocket] = []
                    if pkg not in built_using[s][v][section][pocket]:
                        built_using[s][v][section][pocket].append(
                            (pkg, pmap[rel][pkg]['version']))

    return built_using

built_using_source_format = '%-35s'
built_using_pocket_format = '%-15s'
built_using_component_format = '%-11s'
built_using_usedby_format = '%-35s'

def get_built_using(built_using_map, src):
    out = ""
    src_version = None
    lessthan = False
    if '/' in src:
        src, src_version = src.split('/', 2)
        if src_version.startswith('-'):
            lessthan = True
            src_version = src_version.lstrip('-')
    if src in built_using_map:
        for version in sorted(built_using_map[src]):
            if src_version is not None:
                if lessthan:
                   if apt_pkg.version_compare(version, src_version) >= 0:
                        print("Skipping %s >= %s" % (version, src_version), file=sys.stderr)
                        continue
                elif src_version != version:
                    continue
            for section in sorted(built_using_map[src][version]):
                for pocket in sorted(built_using_map[src][version][section]):
                    for s, v in sorted(
                            built_using_map[src][version][section][pocket]):
                        out += built_using_source_format % ("%s (%s) " % (src, version))
                        out += built_using_pocket_format % pocket
                        out += built_using_component_format % section
                        out += built_using_usedby_format % s
                        out += '\n'

    return out

def get_built_using_header():
    header = built_using_source_format % "Source (version)"
    header += built_using_pocket_format % "Pocket"
    header += built_using_component_format % "Component"
    header += built_using_usedby_format % "Used by"
    header += "\n" + "-" * 78
    return header

def load_ppa(releases=None):
    map = dict()

    if releases is None:
        releases = cve_lib.all_releases

    for rel in releases:
        if '/' not in rel:
            continue
        (base, ppa) = rel.split('/')
        bn = '%s-%s-supported.txt' % (base, ppa)
        fn = os.path.join(os.path.dirname(os.path.dirname(sys.argv[0])), bn)
        #  Fallback to UCT if possible
        if not os.path.isfile(fn) and 'UCT' in os.environ:
            fn = os.path.join(os.environ['UCT'], os.path.basename(fn))

        if not os.path.isfile(fn):
            print("WARN: could not find '%s'. Skipping" % os.path.basename(fn))
            continue

        file = open(fn, "r")
        lines = file.read().split('\n')
        file.close()

        map[rel] = dict()

        pat = re.compile(r'^[a-z]')
        for line in lines:
            if not pat.search(line):
                continue
            # Mock-up an apt Sources file
            if line not in map[rel]:
                map[rel][line] = dict()
                map[rel][line]['pocket'] = ''
                map[rel][line]['section'] = 'main'
    return map

def madison(source, pkg, releases=None):
    answer = dict()
    if not releases:
        releases = cve_lib.releases
    for rel in releases:
        if rel in cve_lib.eol_releases:
            continue
        if '/' in rel:  # we handle ppa overlays differently
            continue
        if pkg in source[rel]:
            name = rel
            if source[rel][pkg]['pocket'] != '':
                name += '-%s' % (source[rel][pkg]['pocket'])
            name += '/%s' % (source[rel][pkg]['section'])
            answer.setdefault(name, dict())
            answer[name].setdefault(pkg, source[rel][pkg]['version'])
    return answer
