def assert_domains_from_rkn(self, domains): self.assertEqual(Session.query(Domain).count(), len(domains)) for d, o in domains.items(): q = Session.query(Domain).filter(Domain.domain==d).one() self.assert_ips(o['ips'], q) self.assert_ips_from_rkn(o['ips']) self.assert_urls(o['urls'], q) self.assert_urls(itertools.chain.from_iterable([i['urls'] for i in domains.values()])) all_ips = set(itertools.chain.from_iterable([i['ips'] for i in domains.values()])) self.assert_ips(all_ips) self.assert_ips_from_rkn(all_ips)
def sync_and_assert(urls): d.sync_urls(urls) Session.commit() self.assert_urls(urls, d) self.assert_urls(urls)
def assert_ips(self, ips, d=None): if d is not None: self.assertEqual(sorted(d.ips_list), sorted(set(ips))) else: self.assertEqual(sorted(i.ip for i in Session.query(IP).all()), sorted(set(ips)))
def assert_ips_from_all(self, ips): for i in ips: ip = Session.query(IP).filter(IP.ip==i).one() print(ip) self.assertTrue(ip.from_rkn) self.assertTrue(ip.from_resolver)
def test_domains_sync_from_resolver(self): domains = { 'test.ru': { 'ips' : [ '1.1.1.1', ], 'urls' : [], }, 'ya.ru': { 'ips' : [ '3.3.3.3', ], 'urls' : [], }, 'r1.com': { 'ips': [ '5.5.5.5', ], 'urls': [], }, } Domain.sync_from_rkn(Session, domains) domains_resolver = { 'test.ru': { 'ips' : [ '1.1.1.1', '2.2.2.2', ], }, 'ya.ru': { 'ips' : [ '8.8.8.8', ], }, 'r1.com': { 'ips': [ '5.5.5.5', ], 'urls': [], }, } Domain.sync_from_resolver(Session, domains_resolver) ips_rkn = [ '3.3.3.3', ] ips_resolver = [ '2.2.2.2', '8.8.8.8', ] ips_rkn_resolver = [ '1.1.1.1', '5.5.5.5', ] self.assert_ips_from_rkn(ips_rkn) self.assert_ips_from_resolver(ips_resolver) self.assert_ips_from_all(ips_rkn_resolver) for i in domains: d = Session.query(Domain).filter(Domain.domain==i).one() self.assertEqual(sorted(d.ips_list), sorted(set(domains[i]['ips']+domains_resolver[i]['ips']))) domains_resolver = { 'test.ru': { 'ips' : [ '1.1.1.1', ], }, 'ya.ru': { 'ips' : [ '8.8.8.8', ], }, 'r1.com': { 'ips': [ '5.5.5.5', '2.2.2.2', ], }, } Domain.sync_from_resolver(Session, domains_resolver) ips_rkn = [ '3.3.3.3', ] ips_resolver = [ '2.2.2.2', '8.8.8.8', ] ips_rkn_resolver = [ '1.1.1.1', '5.5.5.5', ] self.assert_ips_from_rkn(ips_rkn) self.assert_ips_from_resolver(ips_resolver) self.assert_ips_from_all(ips_rkn_resolver) for i in domains: d = Session.query(Domain).filter(Domain.domain==i).one() self.assertEqual(sorted(d.ips_list), sorted(set(domains[i]['ips']+domains_resolver[i]['ips'])))
def assert_urls(self, urls, d=None): if d is not None: self.assertEqual(sorted(urls), sorted(d.urls_list)) else: self.assertEqual(sorted(u.url for u in Session.query(URL).all()), sorted(urls))
def add_domain(self, name='rknfilter.ru'): d = Domain(domain=name) Session.add(d) return d
def tearDown(self): Base.metadata.drop_all(engine) Session.remove()
def sync_and_assert(ips, ips_stored=None): d.sync_ips(ips, ips_stored=ips_stored) Session.commit() self.assert_ips(ips, d)
def __exit__(self, *args, **kwargs): Session.remove()
import logging logging.basicConfig() logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO) #from rknfilter.targets import FetchTarget #a = FetchTarget() #with open('/tmp/dump.xml', 'r') as f: # all_ = f.read() #parsed = a.dump_parser.parse(all_) #a.store_dump(parsed) #from rknfilter.targets import ResolverTarget from rknfilter.db import Session from rknfilter.db import Resource, Domain a = Session.query(Resource).all()[0] print(a.id) #b = Domain(domain='2test') #a[0].domains.append(b) #a[0].domains.append(b) #a[0].domains.append(b) #a[0].domains.append(b) #Session.flush() #Session.commit() #print(Session.query(Resource).join(Resource.domains).filter(Domain.domain=='forcemixes.com').all()[0].urls)